ManishThota commited on
Commit
8323cfd
·
verified ·
1 Parent(s): 7ec133b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -8,7 +8,9 @@ device = 'cuda' if torch.cuda.is_available() else "cpu"
8
  # torch.set_default_device("cuda")
9
 
10
  # Initialize the model and tokenizer
11
- model = AutoModelForCausalLM.from_pretrained("ManishThota/Sparrow").to(device)
 
 
12
  tokenizer = AutoTokenizer.from_pretrained("ManishThota/Sparrow", trust_remote_code=True)
13
 
14
  def predict_answer(image, question):
 
8
  # torch.set_default_device("cuda")
9
 
10
  # Initialize the model and tokenizer
11
+ model = AutoModelForCausalLM.from_pretrained("ManishThota/Sparrow", torch_dtype=torch.float16,
12
+ device_map="auto",
13
+ trust_remote_code=True).to(device)
14
  tokenizer = AutoTokenizer.from_pretrained("ManishThota/Sparrow", trust_remote_code=True)
15
 
16
  def predict_answer(image, question):