Thank you for your reply.
I wanted to get the intermediate 5D output for the original convolution with groups>1.
e.g. Let us use an example similar to your post. Consider an input activation in the shape [batch_size=1, channels=6, height=24, width=24] and a standard conv layer with a weight tensor in the shape [out_channels=6, in_channels=2, height=3, width=3]
We want to get the intermediate 5D outputs of the result of F.conv2d(input, weight, groups=3, stride=1, bias= None, padding= 1) . So, I kept the group loop to loop over the three slices of input and output in accordance with your figure. With mask = identity, the output produced by my function is same as the F.conv2d(input, weight,groups=3, stride=1). However, my function is orders of magnitude slower to run because of the loop.
Can we get the intermediate 5D output with groups=3 without using a loop?