Support this model device-independent. (#52)

- Support this model device-independent. (61ddce6553b8beffc5a859ba31726aab1f9b6979)


Co-authored-by: yejinglai <[email protected]>

speech_conformer_encoder.py

@@ -2477,9 +2477,8 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
             seq_len, batch_size, self.chunk_size, self.left_chunk
         )
 
-        if xs_pad.is_cuda:
-            enc_streaming_mask = enc_streaming_mask.cuda()
-            xs_pad = xs_pad.cuda()
+        if xs_pad.device != "cpu":
+            enc_streaming_mask = enc_streaming_mask.to(xs_pad.device)
 
         input_tensor = xs_pad
         input_tensor, masks = self._forward_embeddings_core(input_tensor, masks)
@@ -2496,8 +2495,8 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
             enc_streaming_mask_nc = self._streaming_mask(
                 seq_len, batch_size, chunk_size_nc, left_chunk_nc
             )
-            if xs_pad.is_cuda:
-                enc_streaming_mask_nc = enc_streaming_mask_nc.cuda()
+            if xs_pad.device != "cpu":
+                enc_streaming_mask_nc = enc_streaming_mask_nc.to(xs_pad.device)
             if masks is not None:
                 hs_mask_nc = masks & enc_streaming_mask_nc
             else: