add the patch

dmlc · Jan 18, 2024 · 85ac333 · 85ac333
1 parent f7e065f
commit 85ac333
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/dgl/graphbolt/feature_fetcher.py b/python/dgl/graphbolt/feature_fetcher.py
@@ -134,6 +134,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 +145,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(