Former-commit-id: 7afbc85daee295cf38dcee9ded5afd87b2c4cfd1
This commit is contained in:
hiyouga 2024-03-24 00:34:54 +08:00
parent 58aa576ae5
commit 75829c8699

View File

@ -110,6 +110,9 @@ def load_model(
if not is_trainable:
model.requires_grad_(False)
model.eval()
for param in model.parameters():
if param.device.type == "cuda":
param.data = param.data.to(model_args.compute_dtype)
else:
model.train()