I think this function is a general way to accomplish what I was saying. I borrowed this code from: https://github.com/pytorch/pytorch/blob/85a7e0f/test/test_distributions.py#L1501 172
def jacobian(inputs, outputs):
    return torch.stack([grad([outputs[:, i].sum()], [inputs], retain_graph=True, create_graph=True)[0]
                        for i in range(outputs.size(1))], dim=-1)