Skip to content

Commit 5c3fe5b

Browse files
[Fixed]RuntimeError: probability tensor contains either inf, nan or element < 0
1 parent d9aab46 commit 5c3fe5b

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

recipes/quickstart/inference/local_inference/multi_modal_infer.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ def load_model_and_processor(model_name: str, hf_token: str):
1515
Load the model and processor based on the 11B or 90B model.
1616
"""
1717
model = MllamaForConditionalGeneration.from_pretrained(model_name, device_map="auto", torch_dtype=torch.bfloat16, token=hf_token)
18+
model = model.bfloat16().cuda()
1819
processor = MllamaProcessor.from_pretrained(model_name, token=hf_token)
1920
return model, processor
2021

0 commit comments

Comments
 (0)