I am facing similar kind of issue but not while creating the layers rather while calling forward.
Below is the code.
class PreActBlock(nn.Module):
    '''Pre-activation version of the BasicBlock.'''
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(PreActBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)

        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False)
            )

    def forward(self,x,planes,w,h):
        c1_mem=c1_spike=torch.zeros(batch_size,planes,w, h)
        c2_mem=c2_spike=torch.zeros(batch_size,planes,w, h)
        
        out = self.conv1(self.bn1(x))
        c1_mem, c1_spike = mem_update(out, c1_mem, c1_spike)
        #shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x
        out = self.conv2(self.bn2(c1_spike))
        c2_mem, c2_spike = mem_update(out, c2_mem, c2_spike)
        out = x+c2_spike 
        return out


out = self.layer1(c_spike,64,32,32)
out = self.layer2(out,128,16,16)
out = self.layer3(out,256,8,8)
out = self.layer4(out,512,4,4)

result = self.forward(*input, **kwargs)
TypeError: forward() takes 2 positional arguments but 5 were given

Could you give me suggestions?
Thank you for your time.