ENSEEIHT/projet-long
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

Merge pull request #109 from ant-Korn/save_to_onnx

Browse source 
Corrected class Up for model to ONNX exporting

Former-commit-id: 060bdcd69886a3082a6f8fb7746e12d5fca3e360
This commit is contained in:
milesial 2019-12-29 20:02:03 +01:00 committed by GitHub
parent c0adecd5d6 70e6ac4258
commit d081192e90
1 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
unet/unet_parts.py
View file

@ -54,8 +54,8 @@ class Up(nn.Module):
def forward(self, x1, x2): def forward(self, x1, x2):
x1 = self.up(x1) x1 = self.up(x1)
# input is CHW # input is CHW
diffY = x2.size()[2] - x1.size()[2] diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
diffX = x2.size()[3] - x1.size()[3] diffX = torch.tensor([x2.size()[3] - x1.size()[3]])
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2]) diffY // 2, diffY - diffY // 2])