So if we wanted the jacobian, we’d have to call torch::autograd::grad right?
I’ve been playing around with what this looks like. This is what I have so far, does this seem correct?
#include <torch/script.h> // One-stop header.
#include <torch/torch.h>

#include <iostream>
#include <memory>

int main() {
  // model with 4 input dimmensions
  // and 10 output classes
  auto model = torch::nn::Linear(4, 10);

  // sample imput
  auto input = torch::ones({1, 4}).requires_grad_(true);

  // forward pass in input
  auto out = model(input);

  auto grad_output = torch::ones({/*batch=*/1, 1}).to(out);

  // compute the jacobian of out wrt input
  for (int i = 0; i < 10; i++) // 10 loops for 10 output classes
  {
    auto gradient = torch::autograd::grad({out.slice(/*dim=*/1, /*start=*/i,
                                                     /*end=*/i+1)},
                                          {input},
                                          /*grad_outputs=*/{grad_output},
                                          /*retain_graph=*/true,
                                          /*create_graph=*/true);
    std::cout << i << '\n';
    std::cout << gradient[0] << '\n';
    // print the gradient of out[:][i] wrt input
  }
  return 0;
}

EDIT*: I’ve changed auto grad_output = torch::ones_like(out); to auto grad_output = torch::ones({/*batch=*/1, 1}).to(out); because I noticed my gradients were off by a factor of 10…