My loss seems to be not decreasing for FCN16 and FCN8, it gets stuck at a constant loss value. Here is the code snippet of my Model. Please let know if the model is correct for FCN16 / FCN8. I am trying to do Image segmentation on biomedical images. I believe my model is correct but still would like to have an additional pair of eyes to verify it. I would also need some suggestion on pretrained features, currently I am not taking any pretrained features.
class FCN(nn.Module):
def init(self, num_classes , mode=‘fcn32’):
super(FCN,self).init()
    self.mode = mode
    
    feats = list(models.vgg16().features.children())
    
    if self.mode == 'fcn16':
        self.feats = nn.Sequential(*feats[0:17])
        self.pool4 = nn.Sequential(*feats[17:24])
        self.pool5 = nn.Sequential(*feats[24:])
        self.res_pool4 = nn.Conv2d(512, num_classes, 1)
    elif self.mode == 'fcn8':
        self.feats = nn.Sequential(*feats[0:10])
        self.pool3 = nn.Sequential(*feats[10:17])
        self.pool4 = nn.Sequential(*feats[17:24])
        self.pool5 = nn.Sequential(*feats[24:])
        self.res_pool3 = nn.Conv2d(256, num_classes, 1)
        self.res_pool4 = nn.Conv2d(512, num_classes, 1)
    else :
        self.feats = nn.Sequential(*feats)
        
    self.fconn = nn.Sequential(nn.Conv2d(512, 4096, 7, padding=3),
                               nn.ReLU(inplace=True),
                               nn.Conv2d(4096, 4096,1),
                               nn.ReLU(inplace=True),
                               nn.Conv2d(4096, num_classes, 1)
                               )
    
    self.activation = nn.Sigmoid()

def forward(self,x):
    
    feats = self.feats(x)
    
   
    
    if self.mode == 'fcn8':
        
        pool3 = self.pool3(feats)
       
        pool4 = self.pool4(pool3)
        
        pool5 = self.pool5(pool4)
        
        fconn = self.fconn(pool5)
        
        
        res_pool3 = self.res_pool3(pool3)
        
        res_pool4 = self.res_pool4(pool4)
        

        # Upsampling to match the shape of 3rd Pooling Layer
        up_pool4 = func.upsample(res_pool4, pool3.size()[2:], mode='bilinear')
        
        # Upsampling to match the shape of 3rd Pooling Layer
        up_fconn = func.upsample(fconn, pool3.size()[2:], mode='bilinear')
        
        
        out = up_pool4 + up_fconn + res_pool3
        
        
    
    elif self.mode == 'fcn16':
         
        pool4 = self.pool4(feats)
        pool5 = self.pool5(pool4)
        
        fconn = self.fconn(pool5)
        res_pool4 = self.res_pool4(pool4)
        
        # Upsampling to match the shape of 4th Pooling Layer
        up_fconn = func.upsample(fconn, pool4.size()[2:], mode='bilinear')
        out = up_fconn + res_pool4
            
    else:
        out = self.fconn(feats)
        
    upsample = func.upsample(out, x.size()[2:],mode='bilinear')
    
    out = self.activation(upsample)
   
    return out