需要注意一点,代码需要写在 backward()
前面,因为钩子函数是在 backward 时执行的。
var.register_hook(lambda grad: torch.where(torch.isnan(grad), torch.ones_like(grad), grad))
total_loss.backward()
需要注意一点,代码需要写在 backward()
前面,因为钩子函数是在 backward 时执行的。
var.register_hook(lambda grad: torch.where(torch.isnan(grad), torch.ones_like(grad), grad))
total_loss.backward()