I think your problem is on conv2 to linear, seen on notebook you have used
 self.fc1 = torch.nn.Linear(16*5*5, 120) which may become incompatible when given 1x224x224 or 1x256x256 images.
For input of [1, 1, 224,224] ,  you should use 16*54*54 which worked.
In your forward method you could print what is output of last conv2-maxpool layer, and tried to initialize dimension with that
    def forward(self, x):
        x = torch.nn.functional.relu(self.conv1(x))  
        x = self.max_pool_1(x) 
        x = torch.nn.functional.relu(self.conv2(x))
        x = self.max_pool_2(x)
        print(x.shape)
        x = x.view(-1, 16*54*54)
        # x = x.view(-1, x.shape[0])
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))
        x = self.fc3(x)

you could always use dummy variable to test it  like below and try to print where is error coming from.
rand = torch.randn([1,1,224,224])
net.forward(rand)

Hope this works.