bugfix: remove optimiser cls before serialising training args

cybershiptrooper · Sep 7, 2024 · 5214f10 · 5214f10
1 parent a370cfd
commit 5214f10
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/iit/utils/io_scripts.py b/iit/utils/io_scripts.py
@@ -36,6 +36,8 @@ def save_model(model_pair: BaseModelPair, args: IOIArgParseNamespace, task: str)
     t.save(ll_model.state_dict(), f"{save_dir}/ll_model_{model_suffix}.pth")
 
     # save training args
+    # remove 'optimizer_cls' key from training args dict
+    training_args.pop("optimizer_cls")
     training_args_file = os.path.join(results_dir, "training_args.json")
     with open(training_args_file, "w") as f:
         json.dump(training_args, f)