I define a customized convolution and the forward propogation is as follows:
class PrunedConv2d(nn.Module):

    def __init__(self, in_channels, out_channels, k_in_mask, k_out_mask, output_width, kernel_size, stride=1, padding=0, bias=True):
        super(PrunedConv2d, self).__init__()
        
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size =  _pair(kernel_size)
        self.stride =  _pair(stride)
        self.padding = _pair(padding)
        

        self.k_in = len(k_in_mask)
        self.k_out = len(k_out_mask)
        self.k_in_mask = Variable(k_in_mask)
        self.k_out_mask = Variable(k_out_mask)
        self.output_width = output_width

        self.weight = nn.Parameter(torch.Tensor(self.k_out, self.k_in, *self.kernel_size))   #32*32*3*3
        self.output = Variable(torch.zeros(1, self.out_channels, self.output_width, self.output_width).cuda())

        if bias:
            self.bias = nn.Parameter(torch.Tensor(self.k_out))
        else:
            self.register_parameter('bias', None)
        
        #count time
        self.squeeze_time = 0
        self.conv_time = 0
        self.expand_time = 0
        
    
    def forward(self, input):
        torch.cuda.synchronize()
        t0 = time.time()
        temp_in = torch.index_select(input, 1, self.k_in_mask)
        torch.cuda.synchronize()
        t1 = time.time()
        self.squeeze_time = t1 -t0
        
        
        torch.cuda.synchronize()
        t0 = time.time()
        temp_out = F.conv2d(temp_in, self.weight, self.bias, self.stride, self.padding)
        torch.cuda.synchronize()
        t1 = time.time()
        self.conv_time = t1 -t0
        
        
        return temp_out

the ‘’‘self.conv_time’’’ is about 0.0015s
but when I change the return of the forwrad function as follows
 return self.output.index_copy_(1, self.k_out_mask, temp_out)

then the ‘’‘self.conv_time’’’ is about 0.0036s
I don’t know what happens , I feels it is quite strange for me.