How to permute the feature maps of each sample in a bathSize?

vision
1

Hi, I’m wondering how to permute the feature maps of each sample in a bathSize?
For example, if the bathSize was set to 32 and there are 128 feature maps, the tensor size might be (32, 128, 40, 40) and each of the samples has its own permutation list, so its tensor size with the whole batchSize is (32, 128).
Does anyone know how to permute the channels for different samples?

2

I think you are looking for torch.gather, with features os size (32, 128, 40, 40) and idx (the permutation for each feature map in the batch) os size (32, 128), this should do the trick:

permuted = torch.gather(features, idx, dim=1)