fix pytorch 1.8.0 compatibility

Author: rusty1s

csrc/cuda/segment_coo_cuda.cu

@@ -277,9 +277,9 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index,
         for (int i = dim + 1; i < out.dim(); i++)
           count = count.unsqueeze(-1);
         if (out.is_floating_point())
-          out.div_(count);
+          out.true_divide_(count);
         else
-          out.div_(count, "floor");
+          out.floor_divide_(count);
       }
     });
   });

csrc/scatter.cpp

@@ -130,9 +130,9 @@ class ScatterMean : public torch::autograd::Function<ScatterMean> {
     count.masked_fill_(count < 1, 1);
     count = broadcast(count, out, dim);
     if (out.is_floating_point())
-      out.div_(count);
+      out.true_divide_(count);
     else
-      out.div_(count, "floor");
+      out.floor_divide_(count);
 
     ctx->save_for_backward({index, count});
     if (optional_out.has_value())

torch_scatter/scatter.py

@@ -52,8 +52,10 @@ def scatter_mean(src: torch.Tensor, index: torch.Tensor, dim: int = -1,
     count = scatter_sum(ones, index, index_dim, None, dim_size)
     count[count < 1] = 1
     count = broadcast(count, out, dim)
-    rounding_mode = None if torch.is_floating_point(out) else 'floor'
-    out.div_(count, rounding_mode=rounding_mode)
+    if out.is_floating_point():
+        out.true_divide_(count)
+    else:
+        out.floor_divide_(count)
     return out