wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update unet.py (#1955)

Browse Source
This commit is contained in:
Nand Dalal
2020-05-26 13:28:17 +00:00
committed by GitHub
parent d0ec11b9d6
commit c967b88fc8
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pl_examples/models/unet.py
+4 -1
View File
@@ -99,7 +99,10 @@ class Up(nn.Module):
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)