fix device setting to allow using accelerator cpu (NVIDIA#8084)

* fix device setting to allow using accelerator cpu

Signed-off-by: Oren Amsalem <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Oren Amsalem <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Sasha Meister <[email protected]>

nemo/collections/asr/modules/wav2vec_modules.py

@@ -32,8 +32,6 @@
 from nemo.core.classes.module import NeuralModule
 from nemo.core.neural_types import AcousticEncodedRepresentation, AudioSignal, LengthsType, NeuralType, SpectrogramType
 
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 
 class TransposeLast(torch.nn.Module):
     """
@@ -341,7 +339,7 @@ def apply_transformer(self, x, padding_mask=None):
     def create_padding_mask(self, length):
         # Broadcast to vectorize creating the padding mask
         max_len = max(length)
-        padding_mask = torch.arange(max_len, device=DEVICE)
+        padding_mask = torch.arange(max_len, device=length.device)
 
         # Switch to binary for transformer, 1 for valid tokens, 0 for padding
         padding_mask = (padding_mask.expand(len(length), max_len) < length.unsqueeze(1)).type(torch.uint8)