Calculating hessian matrix

Is there anyone who can provide me an example how I am able to compute the hessian matrix of the loss(d ^2 loss/d^2 weights) of below network using torch.autograd.functional.hessian() ?

class Net(nn.Module):
def init(self):
super(Net, self).init()
self.main = nn.Sequential(
nn.Linear(input_n,h_n),
Swish(),
nn.Linear(h_n,h_n),
Swish(),
nn.Linear(h_n,1),
)
def forward(self,x):
output = self.main(x)
return output
for epoch in range(epochs):
for batch_idx, (x,y) in enumerate(dataloader):
loss = criterion(x,y)
loss.backward()