Repository: OpenNMT/OpenNMT-py

Issue:
Training fails to start with rotary embedding (Latest OpenNMT-py)
This has never happened before. Did just git clone with latest and got following error when starting training:

```
Traceback (most recent call last):
  File "/ai/onmt/onmt-py-3342r/OpenNMT-py/onmt/utils/distributed.py", line 178, in spawned_train
    process_fn(opt, device_id=device_id)
  File "/ai/onmt/onmt-py-3342r/OpenNMT-py/onmt/train_single.py", line 237, in main
    trainer.train(
  File "/ai/onmt/onmt-py-3342r/OpenNMT-py/onmt/trainer.py", line 319, in train
    self._gradient_accumulation(
  File "/ai/onmt/onmt-py-3342r/OpenNMT-py/onmt/trainer.py", line 532, in _gradient_accumulation
    raise exc
  File "/ai/onmt/onmt-py-3342r/OpenNMT-py/onmt/trainer.py", line 496, in _gradient_accumulation
    model_out, attns = self.model(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/ai/onmt/onmt-py-3342r/OpenNMT-py/onmt/models/model.py", line 262, in forward
    enc_out, enc_final_hs, src_len = self.encoder(src, src_len)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/ai/onmt/onmt-py-3342r/OpenNMT-py/onmt/encoders/transformer.py", line 254, in forward
    enc_out = layer(enc_out, mask)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/ai/onmt/onmt-py-3342r/OpenNMT-py/onmt/encoders/transformer.py", line 114, in forward
    context, _ = self.self_attn(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/ai/onmt/onmt-py-3342r/OpenNMT-py/onmt/modules/multi_headed_attn.py", line 470, in forward
    query, key = apply_rotary_emb(
  File "/ai/onmt/onmt-py-3342r/OpenNMT-py/onmt/modules/multi_headed_attn.py", line 48, in apply_rotary_emb
    query_out = torch.view_as_real(query_ * rope).flatten(3)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cpu!

```


Discussion:
yes I know, I will fix it asap. btw did you use a prompt that was > 2048 tokens?

Resolving patch:
diff --git a/onmt/modules/multi_headed_attn.py b/onmt/modules/multi_headed_attn.py
--- 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
                 )


Test patch: