changes for attribute testing

mosaicml · Aug 13, 2024 · b5d1b0d · b5d1b0d
1 parent f2349d4
commit b5d1b0d
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/llmfoundry/models/hf/hf_causal_lm.py b/llmfoundry/models/hf/hf_causal_lm.py
@@ -234,10 +234,13 @@ def build_inner_model(
                 + 'Please `pip install llm-foundry[gpu]`.',
             )
 
-        assert hasattr(
-            model_cls,
-            'from_pretrained',
-        ), 'HF Model class is not supported, check arguments to function call!'
+        if not (
+            hasattr(model_cls, 'from_pretrained') and
+            hasattr(model_cls, 'from_config')
+        ):
+            raise AttributeError(
+                f'{model_cls=} has missing `from_pretrained` and `from_config` support.',
+            )
 
         # Hugging Face copies the modules into the
         # transformers modules cache. On particular systems, this operation seems to cause contention between