correct device for pos_emb

rwth-i6 · Sep 15, 2024 · 99f3b7d · 99f3b7d
1 parent 9c0fe3f
commit 99f3b7d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/i6_models/parts/conformer/mhsa_rel_pos.py b/i6_models/parts/conformer/mhsa_rel_pos.py
@@ -252,7 +252,7 @@ def _sinusoidal_pe(pos_seq: torch.Tensor, embed_dim: int):
 
         sinusoid_input = torch.outer(pos_seq, inv_freq)
 
-        pos_emb = torch.zeros(pos_seq.shape[0], embed_dim)
+        pos_emb = torch.zeros(pos_seq.shape[0], embed_dim, device=pos_seq.device)
 
         pos_emb[:, 0::2] = sinusoid_input.sin()
         pos_emb[:, 1::2] = sinusoid_input.cos()