remove squeeze(1)

dmlc · Sep 12, 2024 · 6a02473 · 6a02473
1 parent 1fc5f8a
commit 6a02473
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/graphbolt/src/fused_csc_sampling_graph.cc b/graphbolt/src/fused_csc_sampling_graph.cc
@@ -1412,7 +1412,8 @@ 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.contiguous().nonzero();
+  auto positive_probs_indices = probs.nonzero();
+  // returns a N-by-1 tensor, compatible with all data_ptr() usage below
   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)) {