How to apply affine_grid and grid_sample in a depthwise separable fashion?

Consider input with N channels, and I want to apply K Spatial transforms to individual channels, to eventually get NK output channels. How do I go about doing that?
AFAIK, it is not supported currently. Is there a workaround to do this?


You can just apply nn.Conv2d() and set groups=in_channels, and out_channels=NK, it should be what you want.
I assume that N=3 and K=5 here.

input = torch.randn(1,3,10,10)
layer = nn.Conv2d(in_channels=3, out_channels=15, 3, groups=3)
out = layer(input)
# output torch.Size([1, 15, 10, 10])

I want to apply the Spatial Transform, which consists of affine_grid and grid_sample functions, to individual channels, not the Covn2d.