I don't think this is a huge issue, and I don't think we should take care of setting the correct device for the user.
One thing we could perhaps do, is adding
device='auto', which will set
device to the
device of the model parameters, if all parameters live on the same device.