Specify device for codebook indices creation (NVIDIA-NeMo#15405)

Bug: torch.arange always creates a tensor on cpu, but the FSQ might be moved to another device, crashing codes function when called (because of self.decode function: indices//self.dim_base_index...)

Fix: just move to device when creating tensor.

Signed-off-by: Ferdydh <ferdydh.hadiwijaya@tum.de>

Author: Ferdydh

nemo/collections/tts/modules/audio_codec_modules.py

@@ -1256,7 +1256,7 @@ def codes(self):
 
         Note that the codebook entries are implicitly defined by the number of levels.
         """
-        indices = torch.arange(self.codebook_size)
+        indices = torch.arange(self.codebook_size, device=self.dim_base_index.device)
         # [D, B, T]
         indices = rearrange(indices, 'B -> 1 B 1')
         # [B, D, T]