Merge pull request #22 from usernameandme/fix_for_issue_21

Changes for Pytorch 1.5.0 (adapted from mapillary inplace_abn repo)

modules/bn.py

@@ -105,8 +105,9 @@ def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation
         super(InPlaceABN, self).__init__(num_features, eps, momentum, affine, activation, slope)
 
     def forward(self, x):
-        return inplace_abn(x, self.weight, self.bias, self.running_mean, self.running_var,
+        x, _, _ = inplace_abn(x, self.weight, self.bias, self.running_mean, self.running_var,
                            self.training, self.momentum, self.eps, self.activation, self.slope)
+        return x
 
 
 class InPlaceABNSync(ABN):
@@ -115,8 +116,9 @@ class InPlaceABNSync(ABN):
     """
 
     def forward(self, x):
-        return inplace_abn_sync(x, self.weight, self.bias, self.running_mean, self.running_var,
+        x, _, _ =  inplace_abn_sync(x, self.weight, self.bias, self.running_mean, self.running_var,
                                    self.training, self.momentum, self.eps, self.activation, self.slope)
+        return x
 
     def __repr__(self):
         rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \

modules/functions.py

@@ -112,11 +112,12 @@ def forward(ctx, x, weight, bias, running_mean, running_var,
         # Output
         ctx.var = var
         ctx.save_for_backward(x, var, weight, bias)
-        return x
+        ctx.mark_non_differentiable(running_mean, running_var)
+        return x, running_mean, running_var
 
     @staticmethod
     @once_differentiable
-    def backward(ctx, dz):
+    def backward(ctx, dz, _drunning_mean, _drunning_var):
         z, var, weight, bias = ctx.saved_tensors
         dz = dz.contiguous()
 
@@ -200,11 +201,12 @@ def forward(cls, ctx, x, weight, bias, running_mean, running_var,
         # Output
         ctx.var = var
         ctx.save_for_backward(x, var, weight, bias)
-        return x
+        ctx.mark_non_differentiable(running_mean, running_var)
+        return x, running_mean, running_var
 
     @staticmethod
     @once_differentiable
-    def backward(ctx, dz):
+    def backward(ctx, dz, _drunning_mean, _drunning_var):
         z, var, weight, bias = ctx.saved_tensors
         dz = dz.contiguous()