I hope this will help you. The model is a multilayer perceptron with relu as activation function.
import torch
from torch import nn

def gradient(y, x, grad_outputs=None):
    """Compute dy/dx @ grad_outputs"""
    if grad_outputs is None:
        grad_outputs = torch.ones_like(y)
    grad = torch.autograd.grad(y, [x], grad_outputs = grad_outputs, create_graph=True)[0]
    return grad

def jacobian(y, x):
    """Compute dy/dx = dy/dx @ grad_outputs; 
    for grad_outputs in[1, 0, ..., 0], [0, 1, 0, ..., 0], ...., [0, ..., 0, 1]"""
    jac = torch.zeros(y.shape[0], x.shape[0]) 
    for i in range(y.shape[0]):
        grad_outputs = torch.zeros_like(y)
        grad_outputs[i] = 1
        jac[i] = gradient(y, x, grad_outputs = grad_outputs)
    return jac

class Linear(nn.Module):
    """costomized linear layer"""
    def __init__(self, in_features, out_features, bias = True, activation_function = None):
        super(Linear, self).__init__()
        self.linear = nn.Linear(in_features, out_features, bias = bias)
        self.activation_function = activation_function if activation_function else lambda x : x
    
    def forward(self, x):
        return self.activation_function(self.linear(x))

class MLP(nn.Module):
    """Multi-layer perceptron"""
    def __init__(self, in_features, hidden_features, hidden_layers, out_features, activation_function = None):
        super(MLP, self).__init__()
        net = []
        net.append(Linear(in_features, hidden_features, True, activation_function))  
        net += [Linear(hidden_features, hidden_features, True, activation_function) for _ in range(hidden_layers)]  
        net.append(Linear(hidden_features, out_features, True, None)) 
        self.net = nn.Sequential(*net)

    def forward(self, x):
        return self.net(x)

input_dim = 3
output_dim = 2
hidden_dim = 5
n_hidden = 2

model = MLP(in_features = input_dim, 
            hidden_features = hidden_dim, 
            hidden_layers = n_hidden, 
            out_features = output_dim, 
            activation_function = torch.relu)

x = torch.tensor([[1., 2., 0.], [1., 0., -1.]]) # batch of 2 tensors
y = model(x)

for param in model.parameters() :
    print(gradient(y, param))
    print()