Skip to content

Instantly share code, notes, and snippets.

@yulkang
Forked from sbarratt/torch_jacobian.py
Created June 12, 2019 16:02
Show Gist options
  • Save yulkang/4b853f0f137dc983d25dc9f7fb4b5b96 to your computer and use it in GitHub Desktop.
Save yulkang/4b853f0f137dc983d25dc9f7fb4b5b96 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