[PyTorch] How to check which GPU device our data used

When I using PyTorch to train a model, I often use GPU_A to train the model, save model. But if I load the model I saved to test some new data, I always put the new data in a different GPU, we called it GPU_B.

We will get an error message. It is a problem we can solve, of course. For example, I can put the model and new data to the same GPU device (“cuda:0”).

model = model.to('cuda:0')

But what I want to know is, is there any way to directly see which device my data is on? I searched the Internet and found that some people really have this demand, so I have a good record here.

Use “get_device()” to check

Note: This method is only useful for Tensor, and it does not seem to work for Tensor still on the CPU.

import torch

a = torch.tensor([5, 3]).to('cuda:3')
print(a.get_device())

Output:

3

It can be confirmed that this method is feasible.

References

• https://discuss.pytorch.org/t/how-to-know-on-which-gpu-the-tensor-is/31793/4