small fixes

csrc/cuda/grid_cuda.cu

@@ -32,9 +32,9 @@ torch::Tensor grid_cpu(torch::Tensor pos, torch::Tensor size,
   cudaSetDevice(pos.get_device());
 
   if (optional_start.has_value())
-    CHECK_CPU(optional_start.value());
+    CHECK_CUDA(optional_start.value());
   if (optional_start.has_value())
-    CHECK_CPU(optional_start.value());
+    CHECK_CUDA(optional_start.value());
 
   pos = pos.view({pos.size(0), -1}).contiguous();
   size = size.contiguous();

torch_cluster/graclus.py

@@ -36,15 +36,21 @@ def graclus_cluster(row: torch.Tensor, col: torch.Tensor,
     mask = row == col
     row, col = row[mask], col[mask]
 
-    # Randomly shuffle nodes.
     if weight is not None:
-        perm = torch.randperm(row.size(0), device=row.device)
+        weight = weight[mask]
+
+    # Randomly shuffle nodes.
+    if weight is None:
+        perm = torch.randperm(row.size(0), dtype=torch.long, device=row.device)
         row, col = row[perm], col[perm]
 
     # To CSR.
     perm = torch.argsort(row)
     row, col = row[perm], col[perm]
 
+    if weight is not None:
+        weight = weight[perm]
+
     deg = row.new_zeros(num_nodes)
     deg.scatter_add_(0, row, torch.ones_like(row))
     rowptr = row.new_zeros(num_nodes + 1)