From 18ffce36b5ee0809f2e2905c2fd44843a3725ea0 Mon Sep 17 00:00:00 2001 From: hiyouga Date: Sat, 9 Mar 2024 22:37:16 +0800 Subject: [PATCH] fix #2732 --- src/llmtuner/model/loader.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/llmtuner/model/loader.py b/src/llmtuner/model/loader.py index 531d3ea3..0f886c37 100644 --- 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()