fix rope device for long sequence (OpenNMT#2514)

* fix rope device for long sequence * restore device removed by mistake
vince62s · Nov 13, 2023 · fd2f145 · fd2f145
1 parent f3059a5
commit fd2f145
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/onmt/modules/multi_headed_attn.py b/onmt/modules/multi_headed_attn.py
@@ -421,7 +421,7 @@ def forward(
                     if seqlen > self.rope.size(0):
                         self.rope = rotaryembeddings(
                             self.dim_per_head, maxseqlen=(seqlen + 2048)
-                        )
+                        ).to(self.rope.device)
                     rope = self.rope[start_pos : start_pos + seqlen]
                     query, key = apply_rotary_emb(
                         query, key, rope, interleave=self.rotary_interleave
@@ -465,8 +465,8 @@ def forward(
                 if seqlen > self.rope.size(0):
                     self.rope = rotaryembeddings(
                         self.dim_per_head, maxseqlen=(seqlen + 2048)
-                    )
-                rope = self.rope[start_pos : start_pos + seqlen]
+                    ).to(self.rope.device)
+                rope = self.rope[start_pos : start_pos + seqlen].to(query.device)
                 query, key = apply_rotary_emb(
                     query, key, rope, interleave=self.rotary_interleave
                 )