set Resample2dFunctionBackward returned number of gradients to 4

phuongman · Jul 31, 2019 · 7103404 · 7103404
1 parent 39c8153
commit 7103404
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 11 deletions.
diff --git a/models.py b/models.py
@@ -5,15 +5,26 @@
 import math
 import numpy as np
 
-from networks.resample2d_package.resample2d import Resample2d
-from networks.channelnorm_package.channelnorm import ChannelNorm
-
-from networks import FlowNetC
-from networks import FlowNetS
-from networks import FlowNetSD
-from networks import FlowNetFusion
-
-from networks.submodules import *
+try:
+    from networks.resample2d_package.resample2d import Resample2d
+    from networks.channelnorm_package.channelnorm import ChannelNorm
+
+    from networks import FlowNetC
+    from networks import FlowNetS
+    from networks import FlowNetSD
+    from networks import FlowNetFusion
+
+    from networks.submodules import *
+except:
+    from .networks.resample2d_package.resample2d import Resample2d
+    from .networks.channelnorm_package.channelnorm import ChannelNorm
+
+    from .networks import FlowNetC
+    from .networks import FlowNetS
+    from .networks import FlowNetSD
+    from .networks import FlowNetFusion
+
+    from .networks.submodules import *
 'Parameter count = 162,518,834'
 
 class FlowNet2(nn.Module):

diff --git a/networks/resample2d_package/resample2d.py b/networks/resample2d_package/resample2d.py
@@ -33,9 +33,9 @@ def backward(ctx, grad_output):
 
         resample2d_cuda.backward(input1, input2, grad_output.data,
                                  grad_input1.data, grad_input2.data,
-                                 ctx.kernel_size)
+                                 ctx.kernel_size, ctx.bilinear)
 
-        return grad_input1, grad_input2, None
+        return grad_input1, grad_input2, None, None
 
 class Resample2d(Module):