fix for unweighted graph partition

csrc/cpu/metis_cpu.cpp

@@ -7,7 +7,7 @@
 #include "utils.h"
 
 torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,int64_t num_parts,
-                            torch::optional<torch::Tensor> edge_wgt,  bool recursive) {
+                            torch::Tensor edge_wgt,  bool recursive) {
 #ifdef WITH_METIS
   CHECK_CPU(rowptr);
   CHECK_CPU(col);
@@ -17,11 +17,7 @@ torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,int64_t num_
 
   auto *xadj = rowptr.data_ptr<int64_t>();
   auto *adjncy = col.data_ptr<int64_t>();
-  if (edge_wgt==nullptr){
-    adjwgt = nullptr
-  }else{
   auto *adjwgt = edge_wgt.data_ptr<int64_t>();
-  }
 
   int64_t ncon = 1;
   int64_t objval = -1;

csrc/cpu/metis_cpu.h

@@ -3,5 +3,5 @@
 #include <torch/extension.h>
 
 torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
-                            int64_t num_parts,torch::optional<torch::Tensor> edge_wgt,
+                            int64_t num_parts,torch::Tensor edge_wgt,
                             bool recursive);

torch_sparse/metis.py

@@ -7,11 +7,11 @@ from torch_sparse.utils import cartesian1d
 
 
 def metis_wgt(x):
-    if len(x.unique()) == 1:
-        return None
     t1, t2 = cartesian1d(x, x)
     diff = t1 - t2
     diff = diff[diff != 0]
+    if len(diff) == 0:
+        return x.long()
     res = diff.abs().min()
     bod = x.max() - x.min()
     scale = (res / bod).item()