Numpy has apply_along_axis but it doesn’t look like pytorch has that unless I’m missing something.
I’m looking for a way to apply a custom (maybe learned too) convolution onto each feature map individually.
So if you have a (64,32,128,128) tensor you want a say, (5,5) filter to be applied to all 64*32 (128,128) images.
Maybe I’m not thinking in arrays right, what if I reshape it to (64*32,1,128,128) and do convolutions the usual way? Then reshape back? Seems like it should work, would reshaping it naively break things though?