ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Fix Loader issue with dtype fp16 (#10907)

Browse source
This commit is contained in:
Guancheng Fu 2024-04-29 10:16:02 +08:00 committed by GitHub
parent c9fac8c26b
commit fbcd7bc737
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
python/llm/src/ipex_llm/transformers/loader.py
View file

@ -59,6 +59,8 @@ def load_model(
model_kwargs["trust_remote_code"] = True
if low_bit == "bf16":
model_kwargs.update({"load_in_low_bit": low_bit, "torch_dtype": torch.bfloat16})
elif low_bit == "fp16":
model_kwargs.update({"load_in_low_bit": low_bit, "torch_dtype": torch.float16})
else:
model_kwargs.update({"load_in_low_bit": low_bit, "torch_dtype": 'auto'})