As far as I understand you don’t need to backpropagte through the MI calculation. If that’s correct, I would recommend to just use whatever seems simpler. E.g. if sklearn already has some utility functions which make your life easier, just use it and calculate the necessary indices for your feature selection. Note that I also assume this computation will only be applied once in a while, so that the performance doesn’t really matter and thus the missing GPU support in numpy is not a huge bottleneck.


If you would like to select some feature maps from the conv activation, you could simply index them in the forward method of your model.


As long as you calculate the feature indices for each sample in the batch, step 2 should work just fine.