Download classification head if available

h2oai · Nov 6, 2023 · f5e1448 · f5e1448
1 parent b86709a
commit f5e1448
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/llm_studio/app_utils/sections/experiment.py b/llm_studio/app_utils/sections/experiment.py
@@ -1638,6 +1638,7 @@ async def experiment_download_model(q: Q, error: str = ""):
             "config.json",
             "added_tokens.json",
             "model_card.md",
+            "classification_head.pth",
         ]
 
         for file in FILES_TO_PUSH: