Relation between Jacobians and gradients of neural network's forward pass w.r.t. inputs

Question

Let net be an instance of a subclass of PyTorch's nn.Module that implements a 1D CNN whose net.forward method takes as inputs a tensor x with x.shape == torch.Size([1, 1, 224]) representing 224-feature points and returns as output a tensor y_hat with y_hat.shape == torch.Size([1, 3]) representing logits whose argmax is the index of the predicted class for x.

Given a dataset X of input samples with X.shape == torch.Size([194109, 1, 224]), the following code gives me a list of tensors each one of shape torch.Size([1, 3, 1, 1, 224]) storing the Jacobian matrix of the function net.forward evaluated at each input point belonging to the dataset X.

Jacobian_list = []

for i in tqdm(range(X.shape[0])):
    point = X[i:i+1]
    y_hat = net.forward(point)
    M = jacobian(net.forward, point)
    Jacobian_list.append(M)

The above code gives me what I want but it is quite slow. I would like to write a faster version possibly avoiding the for loop and using torch's autograd feature invoking the forward pass only once on the entire dataset X. However I don't know how the implementation of such option would relate to my list of Jacobians.

My end goal is to perform an analysis for feature importance and selection (like saliency maps).

I tried a couple of code snippets but since their results differ, I cannot understand how to interpret them.

First attempt:

input_points = X.requires_grad_(True)
raw_outputs = net(input_points)
class_outputs = torch.argmax(raw_outputs, dim=1).view(-1, 1)
raw_outputs.backward(torch.ones_like(raw_outputs), retain_graph=True)
gradient_list = []
# Backward pass to compute gradients for each class output
for i in range(0, 3):
    # Backward pass for the specific class output
    raw_outputs[:, i].backward(torch.ones_like(raw_outputs[:, i]), retain_graph=True)
    # Store the gradients for the i-th class output
    gradients = input_points.grad.squeeze()
    gradient_list.append(gradients)
    print(f"Gradients w.r.t. output channel {i}: {gradients}")

Second attempt:

input_points = Variable(X, requires_grad=True)
raw_outputs = net(input_points)
class_outputs = torch.argmax(raw_outputs, dim=1).view(-1, 1)
for i in tqdm(range(gradients.shape[0])):
    k = class_outputs[i, 0]
    raw_outputs[i, k].backward(retain_graph=True)
    gradients[i, :] = input_points.grad.data.squeeze()[i, :]
    print(f"Gradients w.r.t. output channel {k}: {gradients}")

How do the above code snippets relate to the list of Jacobians?