Corrected class Up for model to ONNX exporting

Former-commit-id: 2dceff8c909b16530e1291118a2165ea95744afd
This commit is contained in:
Anton Kornilov 2019-12-29 18:37:47 +03:00
parent c0adecd5d6
commit 70e6ac4258

View file

@ -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])