fix: use device-agnostic autocast and tensor placement for ROCm/XPU compatibility

nemo/collections/asr/losses/rnnt_pytorch.py

@@ -182,7 +182,7 @@ def compute_forward_prob(self, acts, duration_acts, labels, act_lens, label_lens
         B, T, U, _ = acts.shape
 
         log_alpha = torch.zeros(B, T, U)
-        log_alpha = log_alpha.cuda()
+        log_alpha = log_alpha.to(device=acts.device)
         for b in range(B):
             for t in range(T):
                 for u in range(U):
@@ -227,7 +227,7 @@ def compute_forward_prob(self, acts, duration_acts, labels, act_lens, label_lens
 
         log_probs = []
         for b in range(B):
-            tt = torch.Tensor([-1000.0]).cuda()[0]
+            tt = torch.tensor(-1000.0, device=acts.device)
 
             # need to loop over all possible ways that blank with different durations contributes to the final loss.
             for n, l in enumerate(self.durations):

nemo/collections/asr/models/confidence_ensemble.py

@@ -86,11 +86,11 @@ def get_filtered_logprobs(hypothesis: Hypothesis, exclude_blank: bool) -> torch.
             filtered_logprobs.append(align_elem[0])
         filtered_logprobs = torch.stack(filtered_logprobs)
         if torch.cuda.is_available():  # by default logprobs are placed on cpu in nemo
-            filtered_logprobs = filtered_logprobs.cuda()
+            filtered_logprobs = filtered_logprobs.to(device=next(self.parameters()).device)
     else:  # CTC
         logprobs = hypothesis.y_sequence
         if torch.cuda.is_available():  # by default logprobs are placed on cpu in nemo
-            logprobs = logprobs.cuda()
+            logprobs = logprobs.to(device=next(self.parameters()).device)
         if exclude_blank:  # filtering blanks
             labels = logprobs.argmax(dim=-1)
             filtered_logprobs = logprobs[labels != logprobs.shape[1] - 1]

nemo/utils/callbacks/cuda_graph.py

@@ -59,7 +59,7 @@ def struct_copy_one(src):
     elif isinstance(src, dict):
         return {k: struct_copy_one(src[k]) for k in src}
     elif isinstance(src, torch.Tensor):
-        return src.clone().detach().cuda()
+        return src.clone().detach().to(src.device)
     else:
         return src
 

nemo/utils/cast_utils.py

@@ -18,14 +18,28 @@
 import torch
 
 
+def _current_device_type() -> str:
+    """Return the current accelerator device type ('cuda', 'xpu', etc.).
+    Falls back to 'cuda' for backwards compatibility."""
+    if torch.cuda.is_available():
+        return 'cuda'
+    try:
+        if torch.xpu.is_available():
+            return 'xpu'
+    except AttributeError:
+        pass
+    return 'cuda'
+
+
+
 def avoid_bfloat16_autocast_context():
     """
     If the current autocast context is bfloat16,
     cast it to float32
     """
 
     if torch.is_autocast_enabled() and torch.get_autocast_gpu_dtype() == torch.bfloat16:
-        return torch.amp.autocast('cuda', dtype=torch.float32)
+        return torch.amp.autocast(_current_device_type(), dtype=torch.float32)
     else:
         return nullcontext()
 
@@ -38,12 +52,12 @@ def avoid_float16_autocast_context():
 
     if torch.is_autocast_enabled() and torch.get_autocast_gpu_dtype() == torch.float16:
         if torch.jit.is_scripting() or torch.jit.is_tracing():
-            return torch.amp.autocast('cuda', dtype=torch.float32)
+            return torch.amp.autocast(_current_device_type(), dtype=torch.float32)
 
         if torch.cuda.is_bf16_supported():
-            return torch.amp.autocast('cuda', dtype=torch.bfloat16)
+            return torch.amp.autocast(_current_device_type(), dtype=torch.bfloat16)
         else:
-            return torch.amp.autocast('cuda', dtype=torch.float32)
+            return torch.amp.autocast(_current_device_type(), dtype=torch.float32)
     else:
         return nullcontext()