diff --git a/transfer_learning_resnet50_custom_data.py b/transfer_learning_resnet50_custom_data.py
index bc2a748..57ffa0c 100644
--- a/transfer_learning_resnet50_custom_data.py
+++ b/transfer_learning_resnet50_custom_data.py
@@ -11,7 +11,7 @@
 from keras.models import Model
 from keras.utils import np_utils
 from sklearn.utils import shuffle
-from sklearn.cross_validation import train_test_split
+from sklearn.model_selection import train_test_split
 
 img_path = 'elephant.jpg'
 img = image.load_img(img_path, target_size=(224, 224))
@@ -162,4 +162,4 @@
 plt.grid(True)
 plt.legend(['train','val'],loc=4)
 #print plt.style.available # use bmh, classic,ggplot for big pictures
-plt.style.use(['classic'])
\ No newline at end of file
+plt.style.use(['classic'])