PyTorch代码调试利器: 自动print每行代码的Tensor信息

使用:

pip install torchsnooper

1. 装饰函数

@torchsnooper.snoop()
def myfunc(mask, x):

2. 没有自定义函数，但是我们可以使用 with 语句来激活 TorchSnooper

with torchsnooper.snoop():
    for _ in range(10):
        optimizer.zero_grad()
        pred = model(x)
        squared_diff = (y - pred) ** 2
        loss = squared_diff.mean()
        print(loss.item())
        loss.backward()
        optimizer.step()

参考:

PyTorch代码调试利器: 自动print每行代码的Tensor信息