diff --git a/models/bayes_cat_dog_classifier.pth b/models/bayes_128px.pth
similarity index 74%
rename from models/bayes_cat_dog_classifier.pth
rename to models/bayes_128px.pth
index ede32b9..8a660a6 100644
Binary files a/models/bayes_cat_dog_classifier.pth and b/models/bayes_128px.pth differ
diff --git a/models/bayes_256px.pth b/models/bayes_256px.pth
new file mode 100644
index 0000000..eecc77a
Binary files /dev/null and b/models/bayes_256px.pth differ
diff --git a/train.py b/train.py
index f61d3c8..32e0621 100644
--- a/train.py
+++ b/train.py
@@ -50,7 +50,7 @@ def train_model(model_name, spinner):
     dataset = TensorDataset(X_tensor, Y_tensor)
     dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
 
-    model = CatDogClassifier()
+    model = CatDogClassifier(img_size=IMG_SIZE)
     model = model.to(DEVICE)
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)