Update pytest_new_process_trace_benchmarks.py

Author: KRRT7

codeflash/benchmarking/pytest_new_process_trace_benchmarks.py

@@ -6,11 +6,11 @@
 trace_file = sys.argv[3]
 # current working directory
 project_root = Path.cwd()
+
 if __name__ == "__main__":
     import pytest
 
-    pytest.register_assert_rewrite("codeflash.benchmarking.codeflash_trace")
-    pytest.register_assert_rewrite("codeflash.benchmarking.plugin.plugin")
+    pytest.register_assert_rewrite("codeflash")
 
     from codeflash.benchmarking.codeflash_trace import codeflash_trace
     from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin