diff --git a/python/pokedex_test.py b/python/pokedex_test.py
index 7cefe5741d7a96e5f820f7b7f1ca27a987139639..011ab2951b78163e7f483a2f91350766749eab0a 100644
--- a/python/pokedex_test.py
+++ b/python/pokedex_test.py
@@ -54,7 +54,7 @@ for i in range(4):
 
     # --- Predict ---
     predictions = model.predict(img_array, verbose=0)
-    probabilities = predictions[0]
+    probabilities = tf.nn.softmax(predictions[0])
     predicted_class_index = np.argmax(probabilities)
     predicted_label = class_names[predicted_class_index]
     confidence = 100 * probabilities[predicted_class_index]