From f73832aca2191f94a6075a41624cf1cda3db0910 Mon Sep 17 00:00:00 2001
From: Younes Belkada <younesbelkada@gmail.com>
Date: Wed, 20 Mar 2024 16:20:29 +0000
Subject: [PATCH] optimize tags

---
 src/axolotl/core/trainer_builder.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py
index b23a8a1245..0e3348df59 100644
--- a/src/axolotl/core/trainer_builder.py
+++ b/src/axolotl/core/trainer_builder.py
@@ -793,6 +793,12 @@ def __init__(self, cfg, model, tokenizer):
         self.model = model
         self.tokenizer = tokenizer
 
+        # in case the model supports tagging, add the axolotl tag.
+        # This makes sure the tag is correctly pushed even if a user calls 
+        # model.push_to_hub instad of  trainer.push_to_hub.
+        if hasattr(model, "add_model_tags"):
+            model.add_model_tags(["axolotl"])
+
     @property
     def model_ref(self):
         return self._model_ref