From 70e6ac425808bdc8e03a2784b2c57379648a9ff6 Mon Sep 17 00:00:00 2001
From: Anton Kornilov <kranton94@mail.ru>
Date: Sun, 29 Dec 2019 18:37:47 +0300
Subject: [PATCH] Corrected class Up for model to ONNX exporting

Former-commit-id: 2dceff8c909b16530e1291118a2165ea95744afd
---
 unet/unet_parts.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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