Thank you for your replying. The mem_update is called in
class neural network(nn.Module):
   def __init__(self, inplanes, planes,  batch_size, stride=1, option='A'):
        super(neural network, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = BatchNorm(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = BatchNorm(planes)
        self.drop2 = nn.Dropout(0.2)
        self.planes = planes
        self.stride = stride
  def forward(self, x, c1_mem, c1_spike, c2_mem, c2_spike):
        out = self.bn1(self.conv1(x))
        c1_mem, c1_spike = mem_update(out, c1_mem, c1_spike)
        out = self.bn2(self.conv2(c1_spike))
        out += self.shortcut(x)
        c2_mem, c2_spike = mem_update(out, c2_mem, c2_spike)
        c2_spike = self.drop2(c2_spike)
        return c2_spike, c1_mem, c1_spike, c2_mem, c2_spike

I have moved the threshold parameter outside of the mem_update, but this goes against the learnable threshold parameter I want to implement.