fix transducer loss inputs devices

Adel-Moumen · Adel-Moumen · commit 2c491a4662b1 · 2022-07-21T18:57:58.000+02:00
diff --git a/speechbrain/nnet/loss/transducer_loss.py b/speechbrain/nnet/loss/transducer_loss.py
@@ -301,6 +301,8 @@ class TransducerLoss(Module):
     The TranducerLoss(nn.Module) use Transducer(autograd.Function)
     to compute the forward-backward loss and gradients.
 
+    Input tensors must be on a cuda device.
+
     Example
     -------
     >>> import torch
@@ -332,11 +334,18 @@ def __init__(self, blank=0, reduction="mean"):
             err_msg += "export NUMBAPRO_NVVM='/usr/local/cuda/nvvm/lib64/libnvvm.so' \n"
             err_msg += "================================ \n"
             err_msg += "If you use conda:\n"
-            err_msg += "conda install numba cudatoolkit=9.0"
+            err_msg += "conda install numba cudatoolkit=XX (XX is your cuda toolkit version)"
             raise ImportError(err_msg)
 
     def forward(self, logits, labels, T, U):
         """Computes the transducer loss."""
         # Transducer.apply function take log_probs tensor.
-        log_probs = logits.log_softmax(-1)
-        return self.loss(log_probs, labels, T, U, self.blank, self.reduction)
+        if logits.device == labels.device == T.device == U.device == "cuda":
+            log_probs = logits.log_softmax(-1)
+            return self.loss(
+                log_probs, labels, T, U, self.blank, self.reduction
+            )
+        else:
+            raise ValueError(
+                f"Found inputs tensors to be on {[logits.device, labels.device, T.device, U.device]} while needed to be on a 'cuda' device to use the transducer loss."
+            )
