Update model.py

model.py

@@ -17,7 +17,6 @@ class PretrainedUNet(nn.Module):
         super().__init__()
         
         self.base_model = resnet34(weights=ResNet34_Weights.IMAGENET1K_V1)
-        # Modify the first convolution layer to accept 6 channels (2x RGB images) instead of 3
         self.base_model.conv1 = nn.Conv2d(6, 64, kernel_size=7, stride=2, padding=3, bias=False)
 
         self.encoder1 = nn.Sequential(self.base_model.conv1, self.base_model.bn1, self.base_model.relu)
@@ -42,18 +41,14 @@ class PretrainedUNet(nn.Module):
         e3 = self.encoder3(e2)
         e4 = self.encoder4(e3)
         b = self.bottleneck(e4)
-        
         d4 = self.upconv4(b)
         d4 = torch.cat([d4, e4], dim=1)
         d4 = self.decoder4(d4)
-        
         d3 = self.upconv3(d4)
         d3 = torch.cat([d3, e3], dim=1)
         d3 = self.decoder3(d3)
-        
         d2 = self.upconv2(d3)
         d2 = torch.cat([d2, e2], dim=1)
         d2 = self.decoder2(d2)
-        
         d1 = self.final_upconv(d2)
         return self.final_conv(d1)