瀏覽代碼

[Fixed]RuntimeError: probability tensor contains either inf, nan or element < 0

himanshushukla12 7 月之前
父節點
當前提交
5c3fe5b477
共有 1 個文件被更改,包括 1 次插入0 次删除
  1. 1 0
      recipes/quickstart/inference/local_inference/multi_modal_infer.py

+ 1 - 0
recipes/quickstart/inference/local_inference/multi_modal_infer.py

@@ -15,6 +15,7 @@ def load_model_and_processor(model_name: str, hf_token: str):
     Load the model and processor based on the 11B or 90B model.
     Load the model and processor based on the 11B or 90B model.
     """
     """
     model = MllamaForConditionalGeneration.from_pretrained(model_name, device_map="auto", torch_dtype=torch.bfloat16, token=hf_token)
     model = MllamaForConditionalGeneration.from_pretrained(model_name, device_map="auto", torch_dtype=torch.bfloat16, token=hf_token)
+    model = model.bfloat16().cuda()
     processor = MllamaProcessor.from_pretrained(model_name, token=hf_token)
     processor = MllamaProcessor.from_pretrained(model_name, token=hf_token)
     return model, processor
     return model, processor