Consider I have an array of matrix M.shape = (N,D,D) where D is the dimension of co-variance matrix and N is the number of co-variance matrices.
I have a method that convert a single co-variance matrix to a correlation matrix defined as
def convert_to_correlation(m):
    inv = torch.inverse(torch.diag(torch.diag(m)).sqrt())
    r = inv@m@inv
    return r

When I have a list of such co-variance matrix I work as
def convert_all_to_correlation(M):
    res = []
    for i in range(len(M)):
        m = convert_to_correlation(M[i])
        res.append(m)
    res = torch.stack(res)
    return res

There are two problems in above approach;
(1) The array could have large first dimension; hence very slow in iterative loop
(2) My M is an output of neural network (hence part of computation graph); I don’t want my approach in convert_to_all_correlation remove them from graph.
Any help is appreciated! Thanks