Within the calc_loss function, there’s a enumerate over x in order to calculate the hessian. So, I take each element in x, then unsqueeze(0) in order to add back in the batch. Maybe this could be an issue? Even though, this is perfectly fine with torch.det
  y = Net(x)
  laplacian = torch.zeros(x.shape[0])
  for i, xi in enumerate(x):
    yi = Net(xi.unsqueeze(0))
    hess = hessian(Net, xi.unsqueeze(0), create_graph=False)
    laplacian[i] = hess.view(x.shape[1], x.shape[1]).diagonal(offset=0).sum()
  loss1 = (laplacian/y)

EDIT: I’ve been messing around with printing out the shapes of the Tensors and I’ve just had an idea. Within my loss function I have loss1 which is
  y = Net(x)
  laplacian = torch.zeros(x.shape[0])
  for i, xi in enumerate(x):
    yi = Net(xi.unsqueeze(0))
    hess = hessian(Net, xi.unsqueeze(0), create_graph=True)
    laplacian[i] = hess.view(x.shape[1], x.shape[1]).diagonal(offset=0).sum()
  loss1 = (laplacian/y) #calculate loss1 (including laplacian term)

each element of laplacian is a single value where y was create in batch. Could this be the source of the problem? Because y is going have the batch dimension involved within its backward pass but each element of laplacian won’t. If I change the code above to,
  loss1 = torch.zeros(x.shape[0])
  for i, xi in enumerate(x):
    yi = Net(xi.unsqueeze(0))
    hess = hessian(Net, xi.unsqueeze(0), create_graph=False)
    lap = hess.view(x.shape[1], x.shape[1]).diagonal(offset=0).sum()
    loss1[i] = lap/yi

and set B to say 10, I can run the code and the shapes seem to be ok. Could this be the problem?