Fix tensor device mismatch issue (facebookresearch#4)

shipmetospace · Aug 29, 2023 · 77a1c95 · 77a1c95
1 parent e09c3a8
commit 77a1c95
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sonar/models/sonar_text/model.py b/sonar/models/sonar_text/model.py
@@ -77,7 +77,7 @@ def sentence_embedding_pooling(
         """
 
         if padding_mask is None:
-            padding_mask = torch.zeros(seqs.shape[:2])
+            padding_mask = torch.zeros(seqs.shape[:2], device=seqs.device)
 
         if pooling == Pooling.LAST:
             seq_length = compute_seq_length(padding_mask, _neg_inf)