update

app.py

@@ -53,7 +53,7 @@ def preprocess_input(input):
 def predict(img):
     x = torch.tensor(preprocess_input(img), dtype=torch.float32)
     with torch.no_grad():
-        return model(x.unsqueeze(0)).argmax().item()
+        return model(x.unsqueeze(0)).argmax(1).item()
     
 # Design UI
 gr.Interface(fn=predict,