Last Updated on 2021-06-28 by Clay
Today, when I load a model I trained with PyTorch and I want to use it to process some new data, I got a error report:
RuntimeError: Input and parameter tensors are not at the same device, found input tensor at cuda:1 and parameter tensor at cuda:0
This error message has some difference with my article title. Because I found that the occurrence rate of this kind of error is relatively high: we load a model which is trained by GPU device, and then we load it to use CPU.
Where does the problem occur? In fact, the error message is straightforward:
Input and parameter tensors are not at the same device.
Most likely the model we load is placed in the GPU’s memory by default, and the new data we want to use for processing is placed in the CPU.
Solution
The solution is very simple: we put the data of different devices on the same device:
Like this:
import torch # Model model = torch.load('MODEL_NAME').cpu()
In this way, the model is moved from the previously preset GPU to the CPU.
By the way, what I actually encountered is that my model loaded in by default on the first GPU, but the data I want to process is on the second GPU.
So in my case, I move the model to the second GPU.
import torch # Model model = torch.load('MODEL_NAME').to('cuda:1')
Then my model works fine.
References
- https://stackoverflow.com/questions/58095627/how-to-fix-input-and-hidden-tensors-are-not-at-the-same-device-in-pytorch
- https://discuss.pytorch.org/t/tensors-are-at-different-cuda-devices/47349
- https://discuss.pytorch.org/t/why-model-to-device-wouldnt-put-tensors-on-a-custom-layer-to-the-same-device/17964/2