Hey guys, I am trying to perform a batched computation of eigenvectors from a tensor of dimension (B, M, M) at each epoch. This computation will be used to compute a term in the loss function of the overall model.
B is the batch size, therefore I want to obtain B separate, individual results.
I am currently achieving this via a simple for loop over the batch dimensions as follows:
batch_eig_vals, batch_eig_vecs = [], []
for i in range(S_B.shape[0]):
    eig_vals, eig_vecs = torch.symeig(S_B[i], eigenvectors=True)
    batch_eig_vals.append(eig_vals)
    batch_eig_vecs.append(eig_vecs)

batch_eig_vals = torch.stack(batch_eig_vals)
batch_eig_vecs = torch.stack(batch_eig_vecs)

However, this is extremely slow: each epoch runs about 5x slower than its counterpart without this particular term.
Are there any clever tricks I can apply to speed this up?