diff --git a/unet/unet_parts.py b/unet/unet_parts.py index 7c64dc1..cf092f3 100644 --- a/unet/unet_parts.py +++ b/unet/unet_parts.py @@ -57,8 +57,8 @@ class Up(nn.Module): def forward(self, x1, x2): x1 = self.up(x1) # input is CHW - diffY = torch.tensor([x2.size()[2] - x1.size()[2]]) - diffX = torch.tensor([x2.size()[3] - x1.size()[3]]) + diffY = x2.size()[2] - x1.size()[2] + diffX = x2.size()[3] - x1.size()[3] x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])