Expected all tensors to be on the same device, but got index is on cpu, different from other tensors on cuda:0

[krishmittal21]

[kaishen-Dotc]

Hi! Thanks for reporting this issue.
I tested with a similar environment and it works fine on my end:
torch: 2.7.1
transformers: 5.0.0rc1
accelerate: 1.7.0

Could you try upgrading your packages to match these versions?
Please let me know if upgrading helps, or share your full environment info so we can investigate further!