Merge pull request #9 from abukhoy/pr-373

trust_remote_code enabled for grok1 only

Author: quic-rishinr

tests/transformers/models/test_causal_lm_models.py

@@ -22,6 +22,7 @@
 from QEfficient.utils.device_utils import get_available_device_id
 from QEfficient.utils.run_utils import ApiRunner
 
+extrenal_models = {"hpcai-tech/grok-1"}
 test_models_qaic = [
     "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
     "gpt2",
@@ -80,10 +81,13 @@ def load_causal_lm_model(model_config):
         num_hidden_layers=model_config["n_layer"],
         attn_implementation="eager",
         low_cpu_mem_usage=False,
-        trust_remote_code=True,
-    )  # Run models for single layers only
+        trust_remote_code=model_config["model_name"] in extrenal_models,
+    )
+    # Convert to FP32 if model is in BF16
+    if getattr(model_hf.config, "torch_dtype", None) == torch.bfloat16:
+        model_hf = model_hf.to(torch.float32)
+
     params = sum(p.numel() for p in model_hf.parameters())
-    model_hf.to(torch.float32)
     model_hf.eval()
     return model_hf, params