Fix example code on README.md

#21

Earlier, I got this error message from the example code:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat2 in method wrapper_CUDA_bmm)

I discovered that the problem originated from the rotary_emb.
To resolve this, I added device_map['language_model.model.rotary_emb'] = 0 to the device map, and it successfully resolved the issue.

boxin-wbx changed pull request status to merged
NVIDIA org

Thank you for your contribution!

Sign up or log in to comment