Solved by cdancette in post #2 
              
              
                Hi, 


You can use tensor.device. It returns you the current device of the tensor, and the index of the GPU if the device is cuda. 


If model and input are not in the same device, it will crash. You cannot do operations between tensors in cpu and tensors in cuda. 


Exemple: 
In [12]: torch.rand(10…