Update unet_model.py

unet_model.py

@@ -13,6 +13,7 @@ class UNet(nn.Module):
                 nn.ReLU(inplace=True),
             )
 
+        # Encoder
         self.enc1 = conv_block(3, 64)
         self.enc2 = conv_block(64, 128)
         self.enc3 = conv_block(128, 256)
@@ -20,8 +21,10 @@ class UNet(nn.Module):
 
         self.pool = nn.MaxPool2d(2)
 
+        # Bottleneck
         self.bottleneck = conv_block(512, 1024)
 
+        # Decoder
         self.upconv4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
         self.dec4 = conv_block(1024, 512)
 
@@ -67,4 +70,4 @@ class UNet(nn.Module):
         u1 = torch.cat([u1, c1], dim=1)
         d1 = self.dec1(u1)
 
-        return torch.sigmoid(self.conv_last(d1))
+        return torch.sigmoid(self.conv_last(d1))  # sigmoid kept (matches BCELoss training)