If your labels are sparse, like: [1, 2, 2, 2] (where 0 is missing), the first solution doesn’t work. The second one works but outputs the mean for label 0, here is my fix:

def mean_by_label(samples, labels):
    ''' select mean(samples), count() from samples group by labels order by labels asc '''
    weight = torch.zeros(labels.max()+1, samples.shape[0]).to(samples.device) # L, N
    weight[labels, torch.arange(samples.shape[0])] = 1
    label_count = weight.sum(dim=1)
    weight = torch.nn.functional.normalize(weight, p=1, dim=1) # l1 normalization
    mean = torch.mm(weight, samples) # L, F
    index = torch.arange(mean.shape[0])[label_count > 0]
    return mean[index], label_count[index]