I am running an MNIST example where I have a Tensor H of shape (batch_size, feat_dim) as output of network and the labels L of shape (batch_size)
I want to create a new tensor of shape (num_classes, feat_dim) by average the rows of H using labels in L
What is the fastest, and elegant way to achieve this in pytorch.