diff --git a/prepare_train_model.py b/prepare_train_model.py
index ddd8bef..bf8667a 100644
--- a/prepare_train_model.py
+++ b/prepare_train_model.py
@@ -12,6 +12,7 @@ import matplotlib.pyplot as plt
 import keras
 from keras.layers import Flatten, Dense
 from keras.applications.vgg16 import VGG16
+from keras.optimizers import SGD
 
 TRAIN_PATH  = 'images/train/'
 TEST_PATH   = 'images/test/'
@@ -92,11 +93,9 @@ input("Press Enter to continue...")
 
 #### Training
 # Compile the model
-# from keras.optimizers import SGD
-
-# model.compile(loss='categorical_crossentropy',
-#               optimizer=SGD(lr=1e-3),
-#               metrics=['accuracy'])
+model.compile(loss='categorical_crossentropy',
+              optimizer=SGD(lr=1e-3),
+              metrics=['accuracy'])
 
 # Start the training process
 # model.fit(x_train, y_train, validation_split=0.30, batch_size=32, epochs=50, verbose=2)