From f5e1448d5727a67ac563075cd55409e4fbe59e81 Mon Sep 17 00:00:00 2001
From: maxjeblick <maxjeblick@users.noreply.github.com>
Date: Mon, 6 Nov 2023 16:30:32 +0100
Subject: [PATCH] Download classification head if available

---
 llm_studio/app_utils/sections/experiment.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/llm_studio/app_utils/sections/experiment.py b/llm_studio/app_utils/sections/experiment.py
index 252fc2bc0..46929f561 100644
--- 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: