feat(up): generic upscale handling

hahnec · Sep 4, 2023 · b1fd4dc · b1fd4dc
1 parent 0b8a24d
commit b1fd4dc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models/unet/sloun_model_adapt.py b/models/unet/sloun_model_adapt.py
@@ -17,7 +17,7 @@ def __init__(self, n_channels, n_classes, bilinear=True):
         self.down3 = (Down(256, 512))
         self.latent_conv = (SlounLatent(512, 512))
         factor = 2 if bilinear else 1
-        self.up1 = (Up(512//2*3, 256 // factor, bilinear))
+        self.up1 = (Up(512//2*3, 256 // factor, bilinear)) if bilinear else (Up(512, 256 // factor, bilinear))
         self.up2 = (Up(256, 128 // factor, bilinear))
         self.up3 = (Up(128, 64, bilinear))
         self.outc = (OutConv(64, n_classes))