I played a bit with your code. You can also feed the full output to the grad function and use a different grad_ouput for each iteration. This seems to be faster for larger sizes, but in practice there shouldn’t be much of a different. Also don’t know if the simple time measurement is a fair comparison.
#include <torch/torch.h>
#include <iostream>
#include <chrono>

// define sizes
unsigned int batch_size = 100;
unsigned int input_dim = 100;
int output_dim = 1000;

torch::Tensor model(const torch::Tensor& input)
{
  unsigned int input_dim = input.sizes().back();
  auto M = torch::linspace(0, 1, input_dim*output_dim).view({input_dim, output_dim});
  return torch::exp(input.matmul(M));
}

int main() {
  // model (uncomment if the function above should be used)
  //auto model = torch::nn::Linear(input_dim, output_dim);

  // sample input
  auto input = torch::linspace(0, 1, batch_size*input_dim).view({batch_size, input_dim}).requires_grad_(true);
  std::cout << "input size: " << input.sizes() << std::endl;

  // forward pass in input
  auto out = model(input);
  std::cout << "output size: " << out.sizes() << std::endl;

  // preallocate jacobian
  auto J = torch::ones({batch_size, output_dim, input_dim});

  // preallocate time variables
  auto start = std::chrono::system_clock::now();
  auto end = std::chrono::system_clock::now();
  auto elapsed = end - start;

  // compute the jacobian of out wrt input
  start = std::chrono::system_clock::now();
  for (int i = 0; i < output_dim; i++)
  {
    auto grad_output = torch::zeros({output_dim});
    grad_output.index_put_({i}, 1);
    grad_output = grad_output.expand({batch_size, output_dim});
    //print(grad_output);
    //std::cout << "grad_output size: " << grad_output.sizes() << std::endl;

    auto gradient = torch::autograd::grad({out},
                                          {input},
                                          /*grad_outputs=*/{grad_output},
                                          /*retain_graph=*/true,
                                          /*create_graph=*/true);
    auto grad = gradient[0].unsqueeze(/*dim=*/1);
    //std::cout << "grad size: " << grad.sizes() << std::endl;
    J.slice(1, i, i+1) = grad;
  }
  end = std::chrono::system_clock::now();
  elapsed = end - start;
  std::cout << "Jacobian: " << std::endl;
  //std::cout << J << std::endl;
  std::cout << "duration: " << elapsed.count() << std::endl;

  // compute the jacobian of out wrt input
  start = std::chrono::system_clock::now();
  for (int i = 0; i < output_dim; i++)
  {
    auto grad_output = torch::ones({1});
    grad_output = grad_output.expand({batch_size, 1});
    //print(grad_output);
    //std::cout << "grad_output size: " << grad_output.sizes() << std::endl;

    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);
    auto grad = gradient[0].unsqueeze(/*dim=*/1);
    //std::cout << "grad size: " << grad.sizes() << std::endl;
    J.slice(1, i, i+1) = grad;
  }
  end = std::chrono::system_clock::now();
  elapsed = end - start;
  std::cout << "Jacobian: " << std::endl;
  //std::cout << J << std::endl;
  std::cout << "duration: " << elapsed.count() << std::endl;

  return 0;
}