Merge branch 'master' into gb_cuda_multigpu_example_refinement

python/dgl/graphbolt/feature_fetcher.py

@@ -92,6 +92,8 @@ def record_stream(tensor):
                     nodes = input_nodes[type_name]
                     if nodes is None:
                         continue
+                    if nodes.is_cuda:
+                        nodes.record_stream(torch.cuda.current_stream())
                     for feature_name in feature_names:
                         node_features[
                             (type_name, feature_name)
@@ -104,6 +106,8 @@ def record_stream(tensor):
                             )
                         )
             else:
+                if input_nodes.is_cuda:
+                    input_nodes.record_stream(torch.cuda.current_stream())
                 for feature_name in self.node_feature_keys:
                     node_features[feature_name] = record_stream(
                         self.feature_store.read(
@@ -134,6 +138,8 @@ def record_stream(tensor):
                         edges = original_edge_ids.get(type_name, None)
                         if edges is None:
                             continue
+                        if edges.is_cuda:
+                            edges.record_stream(torch.cuda.current_stream())
                         for feature_name in feature_names:
                             edge_features[i][
                                 (type_name, feature_name)
@@ -143,6 +149,10 @@ def record_stream(tensor):
                                 )
                             )
                 else:
+                    if original_edge_ids.is_cuda:
+                        original_edge_ids.record_stream(
+                            torch.cuda.current_stream()
+                        )
                     for feature_name in self.edge_feature_keys:
                         edge_features[i][feature_name] = record_stream(
                             self.feature_store.read(