use cuda() for internlm3 gptq

AllentDan · Jan 14, 2025 · be8fc0f · be8fc0f
1 parent af14f54
commit be8fc0f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lmdeploy/lite/apis/gptq.py b/lmdeploy/lite/apis/gptq.py
@@ -98,7 +98,7 @@ def auto_gptq(model: str,
                                                 quantize_config,
                                                 revision=revision,
                                                 torch_dtype=torch_dtype,
-                                                trust_remote_code=True)
+                                                trust_remote_code=True).cuda()
 
     # quantize model, the examples should be list of dict whose keys
     # can only be "input_ids" and "attention_mask"