github
/
InternLM
mirror of https://github.com/InternLM/InternLM
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Update test_hf_model.py

pull/710/head
zhulinJulia24 2024-02-27 12:58:06 +08:00 committed by GitHub
parent 39170d27ef
commit dfc1b09611
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_hf_model.py
View File

@ -125,7 +125,7 @@ class TestMMModel:
# it will be loaded as float32 and might cause OOM Error. # it will be loaded as float32 and might cause OOM Error.
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
model_name, torch_dtype=torch.float16, model_name, torch_dtype=torch.float32,
trust_remote_code=True).cuda() trust_remote_code=True).cuda()
tokenizer = AutoTokenizer.from_pretrained(model_name, tokenizer = AutoTokenizer.from_pretrained(model_name,
trust_remote_code=True) trust_remote_code=True)