Update app.py

app.py

@@ -246,22 +246,27 @@ def inference(image, audio, video, task_type, instruction):
     elif task_type == 'Video Captioning':
         text = ['']  
         text_input = tokenizer(text, padding='longest', return_tensors="pt").to(device) 
-        model = MODEL.load_state_dict(state_dict_video_caption,strict=False)  
+        msg = MODEL.load_state_dict(state_dict_video_caption,strict=False)  
+        model = MODEL
     elif task_type == 'Audio Captioning':
         text = ['']  
         text_input = tokenizer(text, padding='longest', return_tensors="pt").to(device) 
-        model = MODEL.load_state_dict(state_dict_audio_caption,strict=False)  
+        msg = MODEL.load_state_dict(state_dict_audio_caption,strict=False)  
+        model = MODEL
     elif task_type == 'Visual Question Answering':
         question = instruction+'?'+special_answer_token
         text_input = tokenizer(question, padding='longest', return_tensors="pt").to(device) 
-        model = MODEL.load_state_dict(state_dict_vqa,strict=False)  
+        msg = MODEL.load_state_dict(state_dict_vqa,strict=False)  
+        model = MODEL
+        print(msg)
     elif task_type == 'Visual Question Answering':
         question = instruction+'?'+special_answer_token
         text_input = tokenizer(question, padding='longest', return_tensors="pt").to(device) 
-        model = MODEL.load_state_dict(state_dict_video_qa,strict=False)  
+        msg = MODEL.load_state_dict(state_dict_video_qa,strict=False)  
+        model = MODEL
     else:
         raise NotImplemented
-
+    
     if "Video" in task_type:
         image = read_video(image)
     elif "Audio" in task_type: