temporarily add nvtx annotation

rapidsai · Aug 27, 2024 · ffc0e46 · ffc0e46
1 parent a365a24
commit ffc0e46
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/cugraph/cugraph/gnn/data_loading/dist_sampler.py b/python/cugraph/cugraph/gnn/data_loading/dist_sampler.py
@@ -20,6 +20,8 @@
 import cupy
 import cudf
 
+import nvtx
+
 from typing import Union, List, Dict, Tuple, Iterator, Optional
 
 from cugraph.utilities.utils import import_optional, MissingModule
@@ -226,6 +228,7 @@ def get_start_batch_offset(
         else:
             return 0, input_size_is_equal
 
+    @nvtx.annotate("cugraph-buffered-sampling", color="blue")
     def __sample_from_nodes_func(
         self,
         call_id: int,