diff --git a/pokedex.py b/pokedex.py
index 0675508146933be21895c16c3105cdf0c9b8a42a..d60dcaf92942cf36e3470440f64447f4647f1ed9 100644
--- a/pokedex.py
+++ b/pokedex.py
@@ -92,7 +92,7 @@ def simple_xception_netowkr(input_shape, num_classes):
     return keras.Model(inputs, outputs)
 
 
-model = simple_xception_netowkr(input_shape=(256, 256) + (3,), num_classes=156)
+model = simple_xception_netowkr(input_shape=(256, 256) + (3,), num_classes=152)
 
 # Train
 epochs = 25
diff --git a/test.py b/test.py
index 75416206f27979339b46682326952b797280b9ea..4fe02cfb8eff44f8a11f8c449f488d82d246dfa7 100644
--- a/test.py
+++ b/test.py
@@ -16,11 +16,11 @@ temp_ds = keras.utils.image_dataset_from_directory(
 class_names = temp_ds.class_names
 
 # Load model
-model = keras.models.load_model("save_at_3.keras")
+model = keras.models.load_model("save_at_6.keras")
 
 # Load and show image
 img = keras.utils.load_img(
-    "Combined_Dataset/Charmeleon/28d58b5e8c68f76d7986aac99b571377cccac3b6f831fc223ad6123f55fcb001.jpg",
+    "Combined_Dataset/Dracaufeu/0f0537c0761b48be754706eb260cf3634f71238a7cb6961dd39b9914857c6283.jpg",
     target_size=(256, 256)
 )
 plt.imshow(img)