I’m trying to compute Jacobian (and its inverse) of the output of an intermediate layer (block1) with respect to the input to the first layer. The code looks like :
def getInverseJacobian(net2, x):
    # define jacobian matrix
    # x has shape (n_batches X dim of input vector)
    # Take one input point from x and forward it through 1st block 

    jac = torch.zeros(size=(x.shape[1],x.shape[1]))

    y = net2.block1(x)

    for i in range(x.shape[1]):
        jac[i,:] = torch.autograd.grad(y[0][i],x, create_graph=True)[0]



    # Getting inverse of jacobian using Penrose pseudo-inverse
    jac_inverse = torch.pinverse(jac)

    if torch.isnan(jac_inverse).any():
        print('Nan encountered in Jacobian !')
        sys.exit(0)
    
    return jac_inverse

This works well for single data in a batch. How do I convert it to make a Jacobian for complete batch without using loop. This function will be called several times in the training and loop would not be ideal.
Any suggestions? Am I calculating Jacobian efficiently in the first place? [It is accurate though]