Pytorch + Residual Network неожиданно выдает ошибку: TypeError: объект 'NoneType' не вызывается

Я погружаюсь в мир глубокого обучения, и, поскольку Python - мой самый известный язык программирования, я решил начать с фреймворка Pytorch. Следуя руководству, я реализовал ResNet из 50 слоев. Впоследствии я попытался сделать ResNet-18 с более простыми блоками, которые выглядели так:

class baseBlock(torch.nn.Module):
  expansion = 1
  def __init__(self,input_planes,planes,stride=1,dim_change=None):
    super(baseBlock,self).__init__()
    #Declare convolutional layers with batch norms
    self.conv1 = torch.nn.Conv2d(input_planes, planes, stride=stride, kernel_size=3, padding=1)
    self.bn1   = torch.nn.BatchNorm2d(planes)
    self.conv2 = torch.nn.Conv2d(planes, planes, stride=1, kernel_size=3, padding=1)
    self.bn2   = torch.nn.BatchNorm2d(planes)
    self.dim_change = dim_change

  def forward(self,x):
    #Save the residue
    res = x
    output = F.relu(self.bn1(self.conv1(x)))
    output = self.bn2(self.conv2(output))

    if self.dim_change is not None:
      res = self.dim_change(res)

    output += res
    output = F.relu(output)

    return output

Я разделил код на блоки и собственно нейронную сеть, чтобы сделать его более понятным. Остальной код для NN:

class ResNet(torch.nn.Module):
  def __init__(self,block,num_layers,classes=10):
    super(ResNet,self).__init__()
    #according to research paper
    self.input_planes = 64
    self.conv1  = torch.nn.Conv2d(3,64,kernel_size =3,stride=1,padding=1)
    self.bn1    = torch.nn.BatchNorm2d(64)
    self.layer1 = self._layer(block, 64, num_layers[0], stride=1)
    self.layer2 = self._layer(block, 128, num_layers[1], stride=2)
    self.layer3 = self._layer(block, 256, num_layers[2], stride=2)
    self.layer4 = self._layer(block, 512, num_layers[3], stride=2)
    self.avaragePool = torch.nn.AvgPool2d(kernel_size=4,stride=1)
    self.fc = torch.nn.Linear(512*block.expansion, classes)


  def _layer(self,block,planes,num_layers,stride=1):
    dim_change = None
    if stride != 1 or planes != self.input_planes*block.expansion:
      dim_change = torch.nn.Sequential(torch.nn.Conv2d(self.input_planes, planes*block.expansion, kernel_size = 1, stride=stride),
                                                        torch.nn.BatchNorm2d(planes*block.expansion))
      netLayers = []
      netLayers.append(block(self.input_planes,planes,stride=stride, dim_change=dim_change))

      self.input_planes = planes*block.expansion

      for i in range(1,num_layers):
        netLayers.append(block(self.input_planes,planes))
        self.input_planes = planes * block.expansion

      return torch.nn.Sequential(*netLayers)


  def forward(self,x):
    x = F.relu(self.bn1(self.conv1(x)))

    #Problem at this layer
    x = self.layer1(x) <- NoneType object is not callable
    x = self.layer2(x)
    x = self.layer3(x)
    x = self.layer4(x)

    x = F.avg_pool2d(x,4)
    x = x.view(x.size(0),-1)
    x = self.fc(x)

    return x

Проблема возникает в последней функции пересылки в layer1, которая отмечена в коде. Если я прокомментирую это, сеть продолжает вычислять (очевидно, что позже она выдает ошибку, потому что размеры неправильные, но она все равно что-то делает). Если я изменю шаг на этом первом слое, проблема также будет решена, но по тем же причинам размерности все равно выдает ошибку позже. Кажется, я не понимаю, в чем проблема. Я использую стандартный набор данных CIFAR10 для его обучения, и обучающий код настолько стандартен, насколько это возможно.

Если я могу предоставить дополнительную информацию, дайте мне знать. Спасибо