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

View File

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