使用:
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()