fix #2732

2024-03-09 22:37:16 +08:00 · 2024-03-09 22:37:16 +08:00 · 18ffce36b5
parent bdb496644c
commit 18ffce36b5
1 changed files with 4 additions and 1 deletions
--- a/src/llmtuner/model/loader.py
+++ b/src/llmtuner/model/loader.py
@ -109,7 +109,10 @@ def load_model(

    if not is_trainable:
        model.requires_grad_(False)
-        model = model.to(model_args.compute_dtype) if not getattr(model, "quantization_method", None) else model
+        if not getattr(model, "quantization_method", None):
+            for param in filter(lambda p: p.device.type == "cuda", model.parameters()):
+                param.data = param.data.to(model_args.compute_dtype)
+
        model.eval()
    else:
        model.train()