Correctly pass tensor to device
According to the docs, input_tensor.to(device)
doesn't move the tensor to the device but returns a tensor that's on the device.
Fixes GPU predictions
According to the docs, input_tensor.to(device)
doesn't move the tensor to the device but returns a tensor that's on the device.
Fixes GPU predictions