diff --git a/deeplightning/model/unet.py b/deeplightning/model/unet.py index a3e156e..d15e036 100644 --- a/deeplightning/model/unet.py +++ b/deeplightning/model/unet.py @@ -72,6 +72,7 @@ def forward(self, x, y): # shapes. Should we pad from right or left, top or bottom? x = F.pad(x, (0, 1, 0, 1)) x = self.upconv(x) + cropped_y = self.crop(y, x) x = torch.cat((cropped_y, x), dim=1) x = self.conv(x) return x