Merge pull request OpenBMB#153 from whn09/main

combine dtype and device to save CPU memory
lynbklk · May 31, 2024 · fe7184f · fe7184f
2 parents d44bc28 + 5b67e5c
commit fe7184f
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/web_demo_2.5.py b/web_demo_2.5.py
@@ -31,8 +31,7 @@
         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()