skip pylint for deepspeed.utils

axolotl-ai-cloud · Jan 28, 2024 · 5594554 · 5594554
1 parent 27a524c
commit 5594554
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -658,7 +658,9 @@ def load_model(
     skip_prepare_model_for_kbit_training = False
 
     if cfg.model_config_type == "mixtral" and is_deepspeed_zero3_enabled():
-        from deepspeed.utils import set_z3_leaf_modules
+        from deepspeed.utils import (  # pylint: disable=no-name-in-module
+            set_z3_leaf_modules,
+        )
         from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
 
         set_z3_leaf_modules(model, [MixtralSparseMoeBlock])