RuntimeError: Expected x1.dtype() == cos.dtype() to be true, but got false. (Could this error message be improved? If so, please report an enhancement request to PyTorch.)

[JerryDaHeLian]

When I Pre-train LLaMA, there is a error:
File "/usr/local/lib/python3.8/dist-packages/torch/autograd/function.py", line 551, in apply
return super().apply(*args, **kwargs) # type: ignore[misc]
File "/home/xxx/TinyLlama/lit_gpt/fused_rotary_embedding.py", line 39, in forward
rotary_emb.apply_rotary(
RuntimeError: Expected x1.dtype() == cos.dtype() to be true, but got false. (Could this error message be improved? If so, please report an enhancement request to PyTorch.)

Who can help me? 3q!

[Huang Haiduo]

I have solved it! Look at follows:
The bug stem form

lit-llama/lit_llama/model.py

Line 130 in da71ade

 dtype=idx.dtype, 

I managed to change it :
from transformers.utils import is_torch_bf16_gpu_available
dtype=torch.bfloat16 if is_torch_bf16_gpu_available() else torch.float16,