Redesign LossNet for custom dataset

I am sorry, I changed two things

def forward(self, x):
        N, C, H, W = x.size()  #new code                3, 32, 10, 32                                                  
        out = F.relu(self.bn1(self.conv1(x)))
        if (H!=32)|(W!=32):    #new code   <------ I changed here with 2          
            Hpool=H//32                                # 16, 16
            Wpool=W//32
            out=F.max_pool2d(out, (Hpool, Wpool), (Hpool, Wpool))  #end new code
        out = F.relu(self.bn1(self.conv1(x)))       #([10, 64, 32, 32])   out             
        out1 = self.layer1(out)                     #([128, 64, 32, 32])    out1            ([4, 64, 64, 64])
        out2 = self.layer2(out1)                    #([128, 128, 16, 16])   out2            ([4, 128, 32, 32])
        out3 = self.layer3(out2)                    #([128, 256, 6, 6])   out3              ([4, 256, 16, 16])
        out4 = self.layer4(out3)                    #([128, 512, 4, 4])    out4             ([4, 512, 8, 8])    This one should be same

        out = F.avg_pool2d(out4, 8)  #<--------- here I changed 8

The reason for the if statement was to reduce the calculations/runtime if no downsizing was necessary. Since 32, 32 was already working fine.

The H//32 gives back 2 for input size of 64, so the maxpool layer cuts the size in half on that dim.