Update app.py

app.py

@@ -18,7 +18,7 @@ def predict(model_name, image):
     input_tensor = transform(image).unsqueeze(0)
     with torch.no_grad():
         output = model(input_tensor)
-    output_np = torch.softmax(output)[0].numpy()
+    output_np = torch.softmax(output, dim=1)[0].numpy()
     return {clsname: prob for clsname, prob in zip(class_names, output_np)}
  
 interface = gr.Interface(