I have a 2D tensor:
    samples = torch.Tensor([
                     [0.1, 0.1],    #-> group / class 1
                     [0.2, 0.2],    #-> group / class 2
                     [0.4, 0.4],    #-> group / class 2
                     [0.0, 0.0]     #-> group / class 0
              ])

and a label for each sample corresponding to a class:
labels = torch.LongTensor([1, 2, 2, 0])
so len(samples) == len(labels). Now I want to calculate the mean for each class / label. Because there are 3 classes (0, 1 and 2) the final vector should have dimension [n_classes, samples.shape[1]] So the expected solution should be:
    result == torch.Tensor([
                     [0.1, 0.1],
                     [0.3, 0.3], # -> mean of [0.2, 0.2] and [0.4, 0.4]
                     [0.0, 0.0]
              ])

Question: Can this be done in pure pytorch (i.e. no numpy so that I can autograd) and ideally without for loops?