Add softmax to output classes

.gitignore

@@ -0,0 +1,2 @@
+.idea/
+.DS_Store

app.py

@@ -36,12 +36,11 @@ def image_classifier(img_input):
 
     with torch.no_grad():
         pred = model(img)[0]
-        pred -= pred.min()
-        pred /= pred.max()
+        pred = torch.nn.functional.softmax(pred)
         return {classes[i]: float(pred[i]) for i in range(10)}
 
 
 gr.Interface(fn=image_classifier,
              inputs=gr.Image(shape=(28, 28)),
-             outputs="label",
+             outputs=gr.Label(num_top_classes=4),
              examples=["mnist_0.png", "mnist_2.png", "mnist_3.png"]).launch()