From 0b33c5bb74e7b10a12b04d5265d17af582df31db Mon Sep 17 00:00:00 2001
From: pme0 <12113751+pme0@users.noreply.github.com>
Date: Sat, 1 Jul 2023 21:13:53 +0100
Subject: [PATCH] update UNet

---
 deeplightning/model/unet.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/deeplightning/model/unet.py b/deeplightning/model/unet.py
index 8065361..a3e156e 100644
--- a/deeplightning/model/unet.py
+++ b/deeplightning/model/unet.py
@@ -65,13 +65,14 @@ def forward(self, x, y):
         """ x : the input feature map from the expansive path
             y : the corresponding feature map from the contracting path
         """
+        cropped_y = self.crop(y, x)
         x = self.upsample(x)
         # The following padding is not clear from the paper how it should be 
         # implemented but it is necessary in order to get to right feature map 
         # shapes. Should we pad from right or left, top or bottom?
         x = F.pad(x, (0, 1, 0, 1))
         x = self.upconv(x)
-        x = torch.cat((self.crop(y, x), x), dim=1)
+        x = torch.cat((cropped_y, x), dim=1)
         x = self.conv(x)
         return x