Previously, trying to run the model when device is set to cuda:0 gave the following error:

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 index in method wrapper_CUDA__index_select)

This was happening because device was only being used to set the input's device while the model's device was being left as CPU. The fix is to also set the model's device using device

Screenshots

Before Changes [CUDA] (broken):

After Changes [CUDA] (fixed):

After Changes [CPU] (still works):