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()