move out of profiler context

rapidsai · Jul 29, 2024 · 72a36a2 · 72a36a2
1 parent 6a0d22a
commit 72a36a2
Showing 1 changed file with 10 additions and 10 deletions.
diff --git a/python/cudf/cudf/pandas/__main__.py b/python/cudf/cudf/pandas/__main__.py
@@ -73,18 +73,18 @@ def main():
     args = parser.parse_args()
 
     rmm_mode = install()
+    if "managed" in rmm_mode:
+        for key in {
+            "column_view::get_data",
+            "mutable_column_view::get_data",
+            "gather",
+            "hash_join",
+        }:
+            from cudf._lib import pylibcudf
+
+            pylibcudf.experimental.enable_prefetching(key)
     with profile(args.profile, args.line_profile, args.args[0]) as fn:
         args.args[0] = fn
-        if "managed" in rmm_mode:
-            for key in {
-                "column_view::get_data",
-                "mutable_column_view::get_data",
-                "gather",
-                "hash_join",
-            }:
-                from cudf._lib import pylibcudf
-
-                pylibcudf.experimental.enable_prefetching(key)
         if args.module:
             (module,) = args.module
             # run the module passing the remaining arguments