not sure if this is necessary actually

2025-01-16 15:08:34 -05:00
parent a6f2c5d583
commit aad62428e0
1 changed files with 1 additions and 1 deletions
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -405,7 +405,7 @@ class ModelLoader:

        if (
            self.cfg.model_config_type in SUPPORTED_MULTIPACK_MODEL_TYPES
-            and self.cfg.flash_attention
+            and (self.cfg.flash_attention or self.cfg.flex_attention)
            and self.cfg.sample_packing
        ):
            if "auto_map" in self.model_config: