clean up snmg example

rapidsai · May 29, 2024 · 0551718 · 0551718
1 parent b1c9b06
commit 0551718
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_snmg.py b/python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_snmg.py
@@ -106,7 +106,7 @@ def run_train(
     from cugraph_pyg.loader import NeighborLoader
 
     graph_store = GraphStore(is_multi_gpu=True)
-    ixr = torch.tensor_split(data.edge_index, world_size, axis=1)[rank]
+    ixr = torch.tensor_split(data.edge_index, world_size, dim=1)[rank]
     graph_store[
         ("node", "rel", "node"), "coo", False, (data.num_nodes, data.num_nodes)
     ] = ixr
@@ -244,14 +244,14 @@ def run_train(
             print(
                 f"Test Accuracy: {acc_test * 100.0:.4f}%",
             )
-    # dist.barrier()
 
     if rank == 0:
         total_time = round(time.perf_counter() - wall_clock_start, 2)
         print("Total Program Runtime (total_time) =", total_time, "seconds")
         print("total_time - prep_time =", total_time - prep_time, "seconds")
 
     cugraph_comms_shutdown()
+    dist.destroy_process_group()
 
 
 if __name__ == "__main__":