I guess I found the solution.
What I wanted to do was to calculate Laplace’s equation 7.


I found this trick 16 to solve the error.


I was misunderstanding @apaszke 's code 2.


I fixed my code and it seems like what I want. The code is clumsy, but at least it works.
import torch.nn as nn
import torch

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = x ** 3
        return x

class SimpleSumNet(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = x ** 3
        return x.sum(dim=1).unsqueeze(1)

class SigmoidNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 1), nn.Sigmoid())

    def forward(self, x):
        x = self.net(x)
        return x
class ReLUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 1), nn.ReLU())

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

def jacobian(y, x, create_graph = False):
    """
        reference: https://gist.github.com/apaszke/226abdf867c4e9d6698bd198f3b45fb7
    """
    jac = []
    flat_y = y.reshape(-1)
    grad_y = torch.zeros_like(flat_y)
    grad_y = torch.zeros_like(flat_y)
    for i in range(len(flat_y)):
        grad_y[i] = 1.
        grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph)
        grad_x = grad_x.reshape(x.shape)
        jac.append(grad_x.reshape(x.shape))
        grad_y[i] = 0.

    return torch.stack(jac, axis = 0).reshape(y.shape + x.shape)

def laplace(y, x):
    """ Laplace's Equation
        https://en.wikipedia.org/wiki/Laplace%27s_equation
    """
    jac1 = jacobian(y, x, create_graph=True)

    batch_size = jac1.size(0)
    jvs = []
    for i in range(batch_size):
        jvs.append(jac1[i, :, i, :])

    first_dev = torch.stack(jvs, dim=0)

    jac2 = jacobian(first_dev, x)

    jvs = []
    for i in range(batch_size):
        jvs.append(jac2[i, :, :,  i, :])

    second_dev = torch.stack(jvs, dim=0)
    return second_dev.sum(dim = -1).sum(dim = -1)

if __name__=="__main__":

    x = torch.Tensor([[1,1], [3,3], [2,2]]).requires_grad_(True)
    net = SigmoidNet()

    y = net(x)**1 # trick: https://discuss.pytorch.org/t/hessian-of-output-with-respect-to-inputs/60668/10

    print("x:\n",x)
    print("y:\n",y)

    l = laplace(y, x)
    print("laplace:\n", l, l.shape)


I would like to know a more elegant way if there is any.