Fix dtype in AquilaChat2-7B AquilaChat2-7B-16K

Signed-off-by: ldwang <[email protected]>
FlagAI-Open · Oct 20, 2023 · 78671a2 · 78671a2
1 parent 636bf34
commit 78671a2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/flagai/model/aquila2_hf/AquilaChat2-7B-16K.py b/flagai/model/aquila2_hf/AquilaChat2-7B-16K.py
@@ -3,7 +3,7 @@
 device = torch.device("cuda:0")
 model_info = "BAAI/AquilaChat2-7B-16K"
 tokenizer = AutoTokenizer.from_pretrained(model_info, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True, torch_dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True, torch_dtype=torch.half)
 model.eval()
 model.to(device)
 text = "请给出10个要到北京旅游的理由。"

diff --git a/flagai/model/aquila2_hf/AquilaChat2-7B.py b/flagai/model/aquila2_hf/AquilaChat2-7B.py
@@ -3,7 +3,7 @@
 device = torch.device("cuda:0")
 model_info = "BAAI/AquilaChat2-7B"
 tokenizer = AutoTokenizer.from_pretrained(model_info, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True, torch_dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True, torch_dtype=torch.half)
 model.eval()
 model.to(device)
 text = "请给出10个要到北京旅游的理由。"