Skip to content

Instantly share code, notes, and snippets.

@rrmina
Forked from sbarratt/torch_jacobian.py
Created October 2, 2019 11:55
Show Gist options
  • Save rrmina/b61c7045fd6db4f94af56ab87343babd to your computer and use it in GitHub Desktop.
Save rrmina/b61c7045fd6db4f94af56ab87343babd to your computer and use it in GitHub Desktop.
Get the jacobian of a vector-valued function that takes batch inputs, in pytorch.
def get_jacobian(net, x, noutputs):
x = x.squeeze()
n = x.size()[0]
x = x.repeat(noutputs, 1)
x.requires_grad_(True)
y = net(x)
y.backward(torch.eye(noutputs))
return x.grad.data
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment