Update pytest_new_process_trace_benchmarks.py

Author: KRRT7

codeflash/benchmarking/pytest_new_process_trace_benchmarks.py

@@ -1,9 +1,6 @@
 import sys
 from pathlib import Path
 
-from codeflash.benchmarking.codeflash_trace import codeflash_trace
-from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin
-
 benchmarks_root = sys.argv[1]
 tests_root = sys.argv[2]
 trace_file = sys.argv[3]
@@ -12,6 +9,12 @@
 if __name__ == "__main__":
     import pytest
 
+    pytest.register_assert_rewrite("codeflash.benchmarking.codeflash_trace")
+    pytest.register_assert_rewrite("codeflash.benchmarking.plugin.plugin")
+
+    from codeflash.benchmarking.codeflash_trace import codeflash_trace
+    from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin
+
     try:
         # Increase recursion limit to handle recursive benchmarks with tracing overhead
         original_limit = sys.getrecursionlimit()