add ability to define torch_dtype (#697)

casper-hansen · Jan 20, 2025 · ce398ab · ce398ab
1 parent f2cb745
commit ce398ab
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/awq/models/auto.py b/awq/models/auto.py
@@ -1,4 +1,5 @@
 import os
+import torch
 import logging
 from transformers import AutoConfig
 from awq.models import *
@@ -63,6 +64,7 @@ def __init__(self):
     def from_pretrained(
         self,
         model_path,
+        torch_dtype=torch.float16,
         trust_remote_code=True,
         safetensors=True,
         device_map=None,
@@ -78,6 +80,7 @@ def from_pretrained(
         return AWQ_CAUSAL_LM_MODEL_MAP[model_type].from_pretrained(
             model_path,
             model_type,
+            torch_dtype=torch_dtype,
             trust_remote_code=trust_remote_code,
             safetensors=safetensors,
             device_map=device_map,