I have a neural network where the neurons of a single layer (layer 2 in the example) are arranged into some n number of “blocks”. I want to forward propagate through them individually to get to n outputs. Presently, I am doing that using ModuleList and a for loop to go over all the individual blocks. But the code is terribly slow because of this. What is a more efficient way to do this?
Any help would be appreciated, thanks a lot!
class small_net(nn.Module):

    def __init__(self, num_classes=10):
        super(small_net, self).__init__()

        self.bCov1 = BinarizeConv2d(3, 50, kernel_size=3, stride=1, padding=1,bias=True)
        self.bn1 = nn.BatchNorm2d(50, affine=False)
        self.htan1 = nn.Hardtanh(inplace=True)


        self.layer2_neuron_blocks = fixed_block_size(50, 10)
        self.layer2_convs = nn.ModuleList(
                    [BinarizeConv2d(i, 100, kernel_size=3, padding=1, bias=True) for i in self.layer2_neuron_blocks])
        self.mp2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.layer2_bns = nn.ModuleList([nn.BatchNorm2d(100, affine=False) for _ in self.layer2_neuron_blocks])
        
        
        self.fc3 = BinarizeLinear(100 * 16 * 16, 20, bias=True)
        self.bn3 = nn.BatchNorm1d(20, affine=False)
        self.htan3 = nn.Hardtanh(inplace=True)


        self.fc4 = BinarizeLinear(20, num_classes, bias=True)
        self.bn4 = nn.BatchNorm1d(num_classes, affine=False)
        
    def forward(self, x):
        x = self.bCov1(x)
        x = self.bn1(x)     
        x = SignActivation.apply(x)
        
        y2 = [0]+list(np.array(self.layer2_neuron_blocks).cumsum())

        out_layer2 = torch.zeros([x.shape[0], len(self.layer2_neuron_blocks), 100, 16, 16])

        for i in range(len(y2)-1):
            x2 = x[:, y2[i]:y2[i+1], :, :]
            x2 = self.layer2_convs[i](x2)
            x2 = self.mp2(x2)
            x2 = self.layer2_bns[i](x2)
            x2 = SignActivation.apply(x2)
            out_layer2[:, i, :, :, :]  = x2.clone()
        
        x2 = torch.mean(out_layer2, dim=1).to(device)

        x = x2.view(-1, 100 * 16 * 16)

        x = self.fc3(x)
        x = self.bn3(x)
        x = self.htan3(x)

        x = self.fc4(x)
        x = self.bn4(x)

        return x