Update train.py

train.py

@@ -34,14 +34,14 @@ validation_datagen = ImageDataGenerator(rescale=1./255)  # Only rescale for vali
 train_generator = train_datagen.flow_from_directory(
     train_data_dir,
     target_size=(img_width, img_height),
-    batch_size=32,  # Adjust batch size based on GPU memory
+    batch_size=8,  # Adjust batch size based on GPU memory
     class_mode='binary'  # Two classes: cat or dog
 )
 
 validation_generator = validation_datagen.flow_from_directory(
     validation_data_dir,
     target_size=(img_width, img_height),
-    batch_size=32,
+    batch_size=8,
     class_mode='binary'
 )
 
@@ -67,7 +67,7 @@ model.compile(loss='binary_crossentropy',
 # Train the model
 history = model.fit(
     train_generator,
-    epochs=1,  # Adjust number of epochs based on dataset size and validation performance
+    epochs=10,  # Adjust number of epochs based on dataset size and validation performance
     validation_data=validation_generator
 )