diff --git a/scripts/get_e2e_artifacts.py b/scripts/get_e2e_artifacts.py
index 5c6f6f6..5cb6f8e 100755
--- a/scripts/get_e2e_artifacts.py
+++ b/scripts/get_e2e_artifacts.py
@@ -102,14 +102,14 @@
 
 
 def update_path(archive_path: str):
-  """Update the --input_file flag with the new location of the compiled.vmfb"""
+  """Update the --module_file flag with the new location of the compiled.vmfb"""
   backend_path = archive_path.split('traces')[0]  # 'ModuleName/backend_name'.
   compiled_path = os.path.join(FLAGS.artifacts_dir, backend_path,
                                'compiled.vmfb')
   flagfile_path = os.path.join(FLAGS.artifacts_dir, archive_path)
   for line in fileinput.input(files=[flagfile_path], inplace=True):
-    if line.strip().startswith('--input_file'):
-      print(f'--input_file={compiled_path}\n', end='')
+    if line.strip().startswith('--module_file'):
+      print(f'--module_file={compiled_path}\n', end='')
     else:
       print(line, end='')