Update app.py

app.py

@@ -23,7 +23,7 @@ phi_model = AutoModelForCausalLM.from_pretrained(phi_model_name,trust_remote_cod
 # load weights
 model_to_merge = PeftModel.from_pretrained(phi_model,'./model_chkpt/lora_adaptor')
 merged_model   = model_to_merge.merge_and_unload()
-projection.load_state_dict(torch.load('./model_chkpt/step2_projection.pth'))
+projection.load_state_dict(torch.load('./model_chkpt/step2_projection.pth',map_location=torch.device(device)))
 
 def model_generate_ans(img,val_q):
 
@@ -48,7 +48,7 @@ def model_generate_ans(img,val_q):
         phi_output_logits = peft_model(inputs_embeds=val_combined_embeds)['logits'] # 4, 69, 51200
         predicted_word_token_logits = phi_output_logits[:, -1, :].unsqueeze(1) # 4,1,51200
         predicted_word_token = torch.argmax(predicted_word_token_logits, dim = -1) # 4,1
-        predicted_caption[:,g] = predicted_word_token.view(1,-1).to('cpu')
+        predicted_caption[:,g] = predicted_word_token.view(1,-1).to(device)
         
     predicted_captions_decoded = tokenizer.batch_decode(predicted_caption,ignore_index = 50256)