You can use the Hessian vector product from torch.autograd.functional.hvp: torch.autograd.functional.hvp — PyTorch 1.10.0 documentation 3 to estimate the trace (diagonal sum) of the Hessian. Just let the vector be a random +/- vector. This is an unbiased estimator with variance 2(|X|_F^2 - sum_i X_{ii}^2), meaning that if the Hessian is close to diagonal, you’ll need very few samples for an accurate estimate.