Update app.py

app.py

@@ -27,7 +27,7 @@ def predict(inp):
     inp = Image.fromarray(inp.astype('uint8'), 'RGB')
     inp = imgTransforms(inp).unsqueeze(0)
     with torch.no_grad():
-        prediction = torch.nn.functional.softmax(model(inp)[0], dim=0) 
+        prediction = torch.nn.functional.softmax(model(inp)[0]) 
     return {LABELS[i]: float(prediction[i]) for i in range(2)}
     
 examples = [['fiat500.jpg'],['VWUP.jpg']]