update model code

model.py

@@ -137,7 +137,7 @@ def check_file(image_path):
                 outputs = model(images)
                 _, predicted = torch.max(outputs.data, 1)
                 output = predicted
-                return outputs
+                return predicted
 
     def remove_module_from_checkpoint(checkpoint):
         new_state_dict = {}