diff --git a/python/cudf/cudf/pandas/__main__.py b/python/cudf/cudf/pandas/__main__.py index d4cb42d4c0b..591744ce793 100644 --- a/python/cudf/cudf/pandas/__main__.py +++ b/python/cudf/cudf/pandas/__main__.py @@ -73,6 +73,16 @@ 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 args.module: @@ -86,17 +96,6 @@ def main(): sys.argv[:] = args.args runpy.run_path(args.args[0], run_name="__main__") - 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 __name__ == "__main__": main()