Hello, I attempted the solution you proposed here yet I’m still getting the same error. Can you perhaps spot why?
def eval_hessian_diag(loss, model):
	def cat_tensors(tensors):
		t_tensor = torch.empty((sum(t.numel() for t in tensors),))
		offset = 0
		for t in tensors:
			t_tensor[offset:offset + t.numel()] = t.contiguous().view(-1)
			offset += t.numel()
		return t_tensor
	weights = list(model.parameters())
	loss_grad = grad(loss, weights, create_graph=True, retain_graph=True)
	g_tensor = cat_tensors(loss_grad)
	w_tensor = cat_tensors(weights)
    
	l = g_tensor.size(0)
	hessian_diag = torch.empty(l)   
	for idx in range(l):
		grad2rd = grad(g_tensor[idx], w_tensor[idx], retain_graph=True)
		hessian_diag[l] = grad2rd[0]

	return hessian_diag

Just to clarify the difference between my function and the one above is that this function takes in the model and then gets the weights by performing list(model.parameters()). I did this to make it clear where the weights are coming from.