Mean of the vector tensor from the same class in a differentiable way

Hi All,

I have a list of tensors and I have say 500 classes, I want to calculate the mean vector for each class in a differentiable way.

here is my code:
y is the labels
H =torch.randn (800,800)

C = 500
mean_vectors_list = []
mean_vectors_list.append([torch.mean(H[y == cl], dim=0, dtype=torch.float32) for cl in 
range(C)])

The problem is some times H tensor doesn’t have some classes as it’s a randomly selected batch
that gives me Nan.
How to use .scatter_add_ to produce the desired output with cuda