Getting NoneType only on local machine

I am trying to run a code available here: Writing ResNet from Scratch in PyTorch

But when I try to run on my local machine I get a runtime error invalid input with the input being None.

I debugged the code with some prints, and sometimes it prints the shape of the tensor normally, but after it says that type None has no attribute shape, in the same execution.

I tried the code on Google Colab and showed no error, the problem is only on my local machine.

Its like the data is lost in the middle of the execution of the code.

I have no clue what to do to make it work on my local machine.

Could you describe the error a bit more and in particular where in your code you are testing the .shape attribute of which object which then returns None?

In the training loop I check the shape and type of the tensor before being passed to the network and it prints normally

    for epoch in range(num_epochs):
        for i, (images, labels) in enumerate(train_loader):  
            # Move tensors to the configured device
            images = images.to(device)
            labels = labels.to(device)
            
            print(images.shape, images.dtype)

            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            del images, labels, outputs
            torch.cuda.empty_cache()
            gc.collect()

I put some prints inside the forward function of the resnet class and except for the one inside the if None statement prints normally

The error occurs before the call of the print(After...)

class ResNet(nn.Module):
    def __init__(self, block, layers, num_classess = 10):
        super(ResNet, self).__init__()
        self.inplanes = 64
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer0 = self._make_layer(block, 64, layers[0], stride=1)
        self.layer1 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer2 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer3 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        self.fc = nn.Linear(512, num_classess)

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None

        if stride != 1 or self.inplanes != planes:

            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes, kernel_size=1, stride=stride),
                nn.BatchNorm2d(planes)
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.maxpool(x)
        print(f"before layer 0: {x.shape}")
        if x is None:
            print("XNONEX")
        x = self.layer0(x)
        print(f"After: { x.shape }")
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)

        return x

And put some prints inside the resBlock too, that is called by the resnet class

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(ResidualBlock, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_channels)
        )

        self.downsample = downsample
        self.relu = nn.ReLU()
        self.out_channels = out_channels

    def forward(self,x):
        residual = x
        if x is None:
            print("NONE")
        print(f"SHAPE: {x.shape}")
        out = self.conv1(x)
        out = self.conv2(out)
        if self.downsample:
            residual = self.downsample(x)
        out += residual
        # print(f"Out shape = { out.shape }")
        out = self.relu(out)

Here is where it become strange, from nothing the data is lost, and the order of the output of the print is strange too, it prints first the shape and after print it is None but I placed the shape print after the print of if is None

How came the output:

torch.Size([1, 3, 224, 224]) torch.float32
before layer 0: torch.Size([1, 64, 56, 56])
SHAPE: torch.Size([1, 64, 56, 56])
NONE
Traceback (most recent call last):
  File "/home/gilberto/projetos/ResNet/resnet.py", line 216, in <module>
    main()
  File "/home/gilberto/projetos/ResNet/resnet.py", line 185, in main
    outputs = model(images)
  File "/home/gilberto/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/gilberto/projetos/ResNet/resnet.py", line 145, in forward
    x = self.layer0(x)
  File "/home/gilberto/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/gilberto/anaconda3/lib/python3.9/site-packages/torch/nn/modules/container.py", line 204, in forward
    input = module(input)
  File "/home/gilberto/anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/gilberto/projetos/ResNet/resnet.py", line 95, in forward
    print(f"SHAPE: {x.shape}")
AttributeError: 'NoneType' object has no attribute 'shape'