From 1fc5f8a9233a7391cb5ac266452a52cae1328a5c Mon Sep 17 00:00:00 2001
From: Stanley Yang <stanyang@amazon.com>
Date: Tue, 10 Sep 2024 10:11:54 +0000
Subject: [PATCH] change based on #7462

---
 graphbolt/src/fused_csc_sampling_graph.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphbolt/src/fused_csc_sampling_graph.cc b/graphbolt/src/fused_csc_sampling_graph.cc
index 2666e725f2eb..ee2b07c4175d 100644
--- 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)) {