Hi,
I implemented the computation of the Jacobian matrix using the torch.autograd.functional.jacobian  as below. However, it’s still too slow, please help me improve this.
def jacobianBatch(f, wrt):
    '''
    Compute the jacobian (derivaties of outputs w.r.t inputs)
    Input:
            f: pytorch model
            wrt: batch of training data
    Output:
            jacobian: J of the batch of training data
    '''
    jacobian = []
    f.eval()
    for i in range(wrt.shape[0]):
        jac = torch.autograd.functional.jacobian(f, wrt[i])
        jacobian.append(jac)
        #
    return torch.stack(jacobian, 0)