User-agent on CI snapshot download

src/axolotl/integrations/kd/trainer.py

@@ -74,9 +74,6 @@ class AxolotlKDTrainer(AxolotlTrainer):
         target_token_ids_for_loss = target_token_ids[..., 1:, :].contiguous()
         target_mask_for_loss = target_mask[..., 1:, :].contiguous()
 
-        if num_items_in_batch is None:
-            num_items_in_batch = -1
-
         if self.args.kd_zscore_base_temp:
             loss_kd = topk_kd_loss_with_zscore(
                 shift_logits,

tests/conftest.py

@@ -58,11 +58,15 @@ def snapshot_download_w_retry(*args, **kwargs):
     """
     with hf_offline_context(True):
         try:
-            return snapshot_download(*args, **kwargs)
+            return snapshot_download(
+                *args, user_agent={"is_ci": "true", "axolotl": "ci"}, **kwargs
+            )
         except LocalEntryNotFoundError:
             pass
     with hf_offline_context(False):
-        return snapshot_download(*args, **kwargs)
+        return snapshot_download(
+            *args, user_agent={"is_ci": "true", "axolotl": "ci"}, **kwargs
+        )
 
 
 @pytest.fixture(scope="session", autouse=True)