fix model card upload for PEFT models (OpenAccess-AI-Collective#1043)

jinwonkim93 · Jan 5, 2024 · 31d2350 · 31d2350
1 parent f243c21
commit 31d2350
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/axolotl/train.py b/src/axolotl/train.py
@@ -188,6 +188,9 @@ def terminate_handler(_, __, model):
 
  if not cfg.hub_model_id:
  trainer.create_model_card(model_name=cfg.output_dir.lstrip("./"))
+ elif cfg.hub_model_id:
+ # defensively push to the hub to ensure the model card is updated
+ trainer.push_to_hub()
 
  return model, tokenizer