RuntimeError : indices should be either on cpu or on the same device as the indexed tensor (cpu)

[belle9217]

File "/slurmhome/qqq/.local/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 333, in forward
^^^^^^^^^^^^^^^^^
File "/slurmhome/qqq/.local/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 184, in apply_rotary_pos_emb
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
~ ~ ~^ ^ ^^ ^ ^^ ^^ ^ ^ ^^ ^

[belle9217]

how to resolve the bug?thanks for your work.