fix rope device for long sequence (#2514)

* fix rope device for long sequence
* restore device removed by mistake

Author: vince62s

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
                 )