adahessian issue

[zzw-zwzhang]

hv * vi), dim=[2, 3], keepdim=True)) / vi[0, 1].numel() # Hessian diagonal block size is 9 here: torch.sum() reduces the dim 2/3.

IndexError: index 1 is out of bounds for dimension 1 with size 1