Solved by ptrblck in post #2 
              
              
                You could use scatter_add_ and torch.unique to get a similar result. 
However, the result tensor will be sorted according to the class index: 
samples = torch.Tensor([
                     [0.1, 0.1],    #-> group / class 1
                     [0.2, 0.2],    #-> group / class 2
                    …