How to permute output neurons

I am getting outputs from an unsupervised layer with 10 outputs for MNIST. The output neurons are permuted though since we don’t assign them directly. I.e. the 10 outputs are actually: [9., 6., 5., 8., 2., 1., 4., 7., 0., 3.] instead of 0-9.

The shape of the output data is [bs, 10]. How can I permute the 10 so that the layer output is in the right order and maintain the gradient please? Permute seems to only want to permute bs and 10.

Thanks

1 Like

You could index dim1 with your indices as seen here:

x = torch.arange(5 * 10).view(5, 10)
idx = torch.randperm(10)
x = x[:, idx]
print(x)
> tensor([[ 9,  8,  0,  7,  2,  3,  1,  6,  5,  4],
        [19, 18, 10, 17, 12, 13, 11, 16, 15, 14],
        [29, 28, 20, 27, 22, 23, 21, 26, 25, 24],
        [39, 38, 30, 37, 32, 33, 31, 36, 35, 34],
        [49, 48, 40, 47, 42, 43, 41, 46, 45, 44]])

Let me know, if that would work for you.

That should work. Thanks!