Add to(device) to "input_ids"

#18

If the model is taken to the GPU but not the tokenized "input ids" then the following runtime error will be raised:
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu
Therefore, .input_ids should be modified to .input_ids.to(device).

Ref:
https://stackoverflow.com/a/75831568/10216028

VladimirVorobev changed pull request status to merged

Sign up or log in to comment