fixing mode on upsampling parts

Hguimaraes · Dec 11, 2018 · 097160e · 097160e
1 parent f7c81d1
commit 097160e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/nbs/parts.py b/nbs/parts.py
@@ -53,6 +53,7 @@ class UpSamplingBlock(nn.Module):
     def __init__(self, in_ch, out_ch, activation,
                  padding=None, kernel_size=5, mode="linear"):
         super(UpSamplingBlock, self).__init__()
+        self.mode = mode
         self.padding = padding or (kernel_size // 2)
 
         # Convolution block
@@ -100,4 +101,4 @@ def __init__(self, in_ch, out_ch, activation, padding=0):
 
     def forward(self, x, x_enc):
         x = torch.cat([x, x_enc], dim=1)
-        return self.conv(x) - x_enc
+        return self.conv(x) - x_enc