[Examples] The UNet model has some bugs
🐛 Bug

The UNet model definition has some bugs pertaining to bilinear interpolation.
Code sample





  
    
      pytorch-lightning/pl_examples/models/unet.py
    
    
        Lines 35 to 37
      in
      2950f66
    
  
  
    

        
          
           for _ in range(num_layers - 1): 
        

        
          
               layers.append(Up(feats, feats // 2), bilinear) 
        

        
          
               feats //= 2 
        
    
  


In the code above, there seems to be a typo. The bilinear flag should be passed to the function Up(). It has instead been passed to the .append() method of the list.




  
    
      pytorch-lightning/pl_examples/models/unet.py
    
    
        Lines 101 to 104
      in
      2950f66
    
  
  
    

        
          
           if bilinear: 
        

        
          
               self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) 
        

        
          
           else: 
        

        
          
               self.upsample = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2) 
        
    
  


The number of channels once the input passes through either one of these layers is different. For "bilinear", the number of channels remains the same, whereas they decrease to half if a ConvTranspose2d is used. This gives an error in the network's .forward() method.
I wanted to directly use the model for some other application, but not sure how issue 2 should be solved. Maybe use a 1x1 convolution to reduce the channels to half?