Pass trust_remote_code to AutoModelForCausalLM

0090ad81 · Casper Hansen · 637d4abd · 0090ad81
Commit 0090ad81 authored Sep 04, 2023 by Casper Hansen
Hide whitespace changes
Inline Side-by-side

Showing with 7 additions and 1 deletion

awq/models/base.py awq/models/base.py +7 -1

No files found.
--- a/awq/models/base.py
+++ b/awq/models/base.py
@@ -316,7 +316,13 @@ class BaseAWQForCausalLM(nn.Module):
            # Load model weights
            model = AutoModelForCausalLM.from_pretrained(
-                model_filename, device_map=device_map, offload_folder="offload", offload_state_dict=True, torch_dtype=torch_dtype, use_safetensors=safetensors
+                model_filename, 
+                device_map=device_map, 
+                trust_remote_code=trust_remote_code, 
+                offload_folder="offload", 
+                offload_state_dict=True, 
+                torch_dtype=torch_dtype, 
+                use_safetensors=safetensors
            )
            model.eval()