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