change based on #7462

dmlc · Sep 10, 2024 · 1fc5f8a · 1fc5f8a
1 parent bf125d8
commit 1fc5f8a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/graphbolt/src/fused_csc_sampling_graph.cc b/graphbolt/src/fused_csc_sampling_graph.cc
@@ -1412,7 +1412,7 @@ inline int64_t UniformPick(
 /** @brief An operator to perform non-uniform sampling. */
 static torch::Tensor NonUniformPickOp(
     torch::Tensor probs, int64_t fanout, bool replace) {
-  auto positive_probs_indices = probs.nonzero().squeeze(1);
+  auto positive_probs_indices = probs.contiguous().nonzero();
   auto num_positive_probs = positive_probs_indices.size(0);
   if (num_positive_probs == 0) return torch::empty({0}, torch::kLong);
   if ((fanout == -1) || (num_positive_probs <= fanout && !replace)) {