pytorch - 如何对张量 \ 模型的设备 (cpu \ gpu) 设置进行故障排除

Question

I have a torch model that i'm trying to port from CPU do be device independent.

setting the device parameter when creating tensors, or calling model.to(device) to move a full model to the target device solves part of the problem, but there are some "left behind" tensors (like variables created during the forward call)

is there a way to detect these without using an interactive debugger? something like tracing tensor creation to allow discover of tensors that are created on the wrong device?

Answer 1

You could check the garbage collector:

import gc
import torch

s = torch.tensor([2], device='cuda:0')
t = torch.tensor([1])
for obj in gc.get_objects():
    if torch.is_tensor(obj):
        print(obj)

Output:

tensor([2], device='cuda:0')
tensor([1])