From 5b67e5c822260ecea53043344686344aa38b547c Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Mon, 27 May 2024 05:35:21 +0000 Subject: [PATCH] combine dtype and device to save CPU memory --- web_demo_2.5.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/web_demo_2.5.py b/web_demo_2.5.py index cc80a82..b2a4ef1 100644 --- a/web_demo_2.5.py +++ b/web_demo_2.5.py @@ -31,8 +31,7 @@ if 'int4' in model_path: exit() model = AutoModel.from_pretrained(model_path, trust_remote_code=True) else: - model = AutoModel.from_pretrained(model_path, trust_remote_code=True).to(dtype=torch.float16) - model = model.to(device=device) + model = AutoModel.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.float16, device_map=device) tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) model.eval()