fix backward for max min

Author: rusty1s

setup.py

@@ -20,7 +20,7 @@
                       ['cuda/scatter.cpp', 'cuda/scatter_kernel.cu'])
     ]
 
-__version__ = '1.3.0'
+__version__ = '1.3.1'
 url = 'https://github.com/rusty1s/pytorch_scatter'
 
 install_requires = []

torch_scatter/__init__.py

@@ -7,7 +7,7 @@
 from .max import scatter_max
 from .min import scatter_min
 
-__version__ = '1.3.0'
+__version__ = '1.3.1'
 
 __all__ = [
     'scatter_add',

torch_scatter/max.py

@@ -24,8 +24,11 @@ def backward(ctx, grad_out, grad_arg):
 
         grad_src = None
         if ctx.needs_input_grad[1]:
-            grad_src = grad_out.new_zeros(index.size())
-            grad_src.scatter_(ctx.dim, arg.detach(), grad_out)
+            size = list(index.size())
+            size[ctx.dim] += 1
+            grad_src = grad_out.new_zeros(size)
+            grad_src.scatter_(ctx.dim, arg.detach() + 1, grad_out)
+            grad_src = grad_src.narrow(ctx.dim, 1, index.size(ctx.dim))
 
         return None, grad_src, None, None
 

torch_scatter/min.py

@@ -24,8 +24,11 @@ def backward(ctx, grad_out, grad_arg):
 
         grad_src = None
         if ctx.needs_input_grad[1]:
-            grad_src = grad_out.new_zeros(index.size())
-            grad_src.scatter_(ctx.dim, arg.detach(), grad_out)
+            size = list(index.size())
+            size[ctx.dim] += 1
+            grad_src = grad_out.new_zeros(size)
+            grad_src.scatter_(ctx.dim, arg.detach() + 1, grad_out)
+            grad_src = grad_src.narrow(ctx.dim, 1, index.size(ctx.dim))
 
         return None, grad_src, None, None