record stream in both branches

Author: mfbalin

python/dgl/graphbolt/impl/neighbor_sampler.py

@@ -398,6 +398,12 @@ def record_stream(tensor):
                         orig_edge_ids is not None
                         and subgraph.original_edge_ids[etype] is None
                     ):
+                        edge_ids = (
+                            subgraph._edge_ids_in_fused_csc_sampling_graph[
+                                etype
+                            ]
+                        )
+                        edge_ids.record_stream(torch.cuda.current_stream())
                         subgraph.original_edge_ids[etype] = record_stream(
                             index_select(orig_edge_ids, edge_ids)
                         )
@@ -418,6 +424,9 @@ def record_stream(tensor):
                     orig_edge_ids is not None
                     and subgraph.original_edge_ids is None
                 ):
+                    subgraph._edge_ids_in_fused_csc_sampling_graph.record_stream(
+                        torch.cuda.current_stream()
+                    )
                     subgraph.original_edge_ids = record_stream(
                         index_select(
                             orig_edge_ids,