Fix for unit test

rapidsai · Feb 23, 2024 · b04cfc1 · b04cfc1
1 parent cc633b9
commit b04cfc1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/pylibwholegraph/pylibwholegraph/test_utils/test_comm.py b/python/pylibwholegraph/pylibwholegraph/test_utils/test_comm.py
@@ -103,9 +103,9 @@ def host_sample_all_neighbors(
         output_id = output_sample_offset_tensor[i]
         for j in range(end - start):
             output_dest_tensor[output_id + j] = host_csr_col_ptr[start + j]
-            output_center_localid_tensor[output_id + j] = node_id
+            output_center_localid_tensor[output_id + j] = i
             output_edge_gid_tensor[output_id + j] = start + j
-    return output_dest_tensor, output_center_localid_tensor, output_edge_gid_tensor
+    return output_sample_offset_tensor, output_dest_tensor, output_center_localid_tensor, output_edge_gid_tensor
 
 
 def copy_host_1D_tensor_to_wholememory(