?

app.py

@@ -15,6 +15,7 @@ print("Done")
 
 
 def create_medusa_heads(model_id: str):
+    parser = distributed_run.get_args_parser()
     training_args = [
         "--model_name_or_path", model_id,
         "--data_path", "data/ShareGPT_V4.3_unfiltered_cleaned_split.json",
@@ -37,7 +38,10 @@ def create_medusa_heads(model_id: str):
         "--medusa_num_heads", "3",
         "--medusa_num_layers", "1",
     ]
-    distributed_run.run_script_path("medusa/medusa/train/train.py", *training_args)
+    args = parser.parse_args(
+        ["training_script", "medusa/medusa/train/train.py", "training_script_args"] + training_args
+    )
+    distributed_run.run(args)
 
     # Upload the medusa heads to the Hub
     repo_id = f"medusa-{model_id}"