diff --git a/__pycache__/BN.cpython-312.pyc b/__pycache__/BN.cpython-312.pyc
index 66a63a7..59b8f8d 100644
Binary files a/__pycache__/BN.cpython-312.pyc and b/__pycache__/BN.cpython-312.pyc differ
diff --git a/bayes_cat_dog_classifier.pth b/bayes_cat_dog_classifier.pth
new file mode 100644
index 0000000..ede32b9
Binary files /dev/null and b/bayes_cat_dog_classifier.pth differ
diff --git a/dataset/test_set/XD/IMG_2763.jpg b/dataset/test_set/XD/IMG_2763.jpg
deleted file mode 100644
index 5328f87..0000000
Binary files a/dataset/test_set/XD/IMG_2763.jpg and /dev/null differ
diff --git a/dataset/test_set/XD/IMG_3077.jpg b/dataset/test_set/XD/IMG_3077.jpg
deleted file mode 100644
index 9df0ea0..0000000
Binary files a/dataset/test_set/XD/IMG_3077.jpg and /dev/null differ
diff --git a/main.py b/main.py
index ad60870..5517bba 100644
--- a/main.py
+++ b/main.py
@@ -3,9 +3,8 @@ import cv2
 import torch
 import numpy as np
 
-model = torch.load("models/bayes_cat_dog_classifier.pth")
+model = torch.load("bayes_cat_dog_classifier.pth")
 model.eval()
-
 model.to("cuda")
 
 IMG_SIZE = 128
@@ -24,11 +23,22 @@ def predict_image(image_path):
     return "Dog" if predicted == 1 else "Cat"
 
 
+# Cats
 preds = []
-for filename in os.listdir("dataset/test_set/XD/"):
-    img_path = os.path.join("dataset/test_set/XD/", filename)
+for filename in os.listdir("dataset/test_set/cats/"):
+    img_path = os.path.join("dataset/test_set/cats/", filename)
     prediction = predict_image(img_path)
     preds.append(prediction)
 
 print(preds.count("Cat"))
+print(preds.count("Cat") / 1000)
+
+# Dogs
+preds = []
+for filename in os.listdir("dataset/test_set/dogs/"):
+    img_path = os.path.join("dataset/test_set/dogs/", filename)
+    prediction = predict_image(img_path)
+    preds.append(prediction)
+
 print(preds.count("Dog"))
+print(preds.count("Dog") / 1000)
\ No newline at end of file
diff --git a/train.py b/train.py
index d3acdb3..5387bb6 100644
--- a/train.py
+++ b/train.py
@@ -7,8 +7,6 @@ from torch import nn, optim
 from BN import CatDogClassifier
 import time
 
-from main import total_loss, outputs
-
 IMG_SIZE = 128
 
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -28,6 +26,7 @@ def load_images_from_folder(folder, label):
     return data
 
 if __name__ == "__main__":
+
     # Loading the dataset
     cat_data = load_images_from_folder("dataset/training_set/cats", label=0)
     dog_data = load_images_from_folder("dataset/training_set/dogs", label=1)
@@ -51,9 +50,10 @@ if __name__ == "__main__":
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)
 
-    num_epochs = 25
+    num_epochs = 20
     start_time = time.time()
 
+    model.train()
     for epoch in range(num_epochs):
         total_loss = 0
         for images, labels in dataloader:
@@ -67,4 +67,5 @@ if __name__ == "__main__":
         print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(dataloader):.4f}")
 
     print(f"Time taken: {(time.time() - start_time):.2f} seconds")
-    torch.save(model, f"models/bayes_cat_dog_classifier.pth")
\ No newline at end of file
+
+    torch.save(model, f"bayes_cat_dog_classifier.pth")
\ No newline at end of file