Hi, I am seeing an issue on the backward pass when using torch.linalg.eigh on a hermitian matrix with repeated eigenvalues. I was wondering if there is any way to obtain the eigenvector associated with the minimum eigenvalue without the gradients in the backward pass going to nan. I am performing this calculation as a part of the loss function and there are no learnable parameters after this point.
I have set up the following testcase that illustrates the issue I am having, with the gradients on the torch.linalg.eigh calculation causing issues. I am using pytorch version 1.9
import torch

matrix = torch.tensor([[
        [
            1.2039, 0, 0.3690, 0
        ],
        [
            0, -1.2039, 0, 0.3690
        ],
        [
            0.3690, 0, -1.2039, 0
        ],
        [
            0, 0.3690, 0, 1.2039
        ],
    ]], requires_grad=True)

e, v = torch.linalg.eigh(matrix)
min_index = torch.argmin(e)
min_eigenvector = v[:, min_index]
min_state = min_eigenvector
min_sum = torch.sum(min_state) # Sum to obtain a scalar for .backward()
min_sum.backward()
assert not matrix.grad.isnan().any()