Update unet.py (Lightning-AI#1955)

cuent · May 26, 2020 · c967b88 · c967b88
1 parent d0ec11b
commit c967b88
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pl_examples/models/unet.py b/pl_examples/models/unet.py
@@ -99,7 +99,10 @@ def __init__(self, in_ch: int, out_ch: int, bilinear: bool = False):
         super().__init__()
         self.upsample = None
         if bilinear:
-            self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.upsample = nn.Sequential(
+                nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
+                nn.Conv2d(in_ch, in_ch // 2, kernel_size=1),
+            )
         else:
             self.upsample = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2)