In my work, I need to split filter and input feature map by 4 along channel depth and do convolution individualy.
I want to ask how can I do to facilitate this procedure( maybe can get rid of  “for loop”)
I have tried

depth-wise convolution, but I can not reshape weight to correct order to ouput correct result.
torch.multiprocessing, but it seems to be applicable only on those tensor without gradient?

My question is how to get every temp_output faster~thanks
code like this
def forward(self, input, weight, padding = 0, stride = 1):
        number = 4
        number_group = int(weight.size(1)/number)
        temp_output_group = [None]*number_group
        for i in range():
            channel_start = i*number
            channel_end = (i+1)*number
            temp_output = F.conv2d(input[:,channel_start:channel_end], weight[:,channel_start:channel_end], padding=padding, stride=stride)
            temp_output_group.append(temp_output)
        #do something with temp_output group
        return ...