Solved by ptrblck in post #2 
              
              
                I’m not sure, if I understand the use case correctly, but would scattering the output logits using a mapping work? 
Here is a small example: 
batch_size = 10
nb_classes = 2
nb_features = 5
nb_out = 10

x = torch.randn(batch_size, nb_features)
target = torch.randint(0, nb_classes, (batch_size,))

map…