diff --git a/dlutils/pytorch/jacobian.py b/dlutils/pytorch/jacobian.py index 56ae282471097cc18408e9e22583c32e798e4879..c44864b3d5d7c10124f299088a762083fa6dda84 100644 --- a/dlutils/pytorch/jacobian.py +++ b/dlutils/pytorch/jacobian.py @@ -1,5 +1,4 @@ import torch -from torch.autograd.gradcheck import zero_gradients def jacobian(inputs, output): @@ -19,7 +18,7 @@ def jacobian(inputs, output): J = J.cuda() for i in range(num_classes): - zero_gradients(inputs) + inputs.zero_grad() grad_output.zero_() grad_output[:, i] = 1 output.backward(grad_output, retain_graph=True)