diff --git a/demo/guide-python/sklearn_examples.py b/demo/guide-python/sklearn_examples.py
index a5a91ffa1..3908e0333 100755
--- a/demo/guide-python/sklearn_examples.py
+++ b/demo/guide-python/sklearn_examples.py
@@ -8,12 +8,8 @@ import pickle
 import xgboost as xgb
 
 import numpy as np
-try:
-    from sklearn.model_selection import KFold, train_test_split
-except:
-    from sklearn.cross_validation import KFold, train_test_split
+from sklearn.model_selection import KFold, train_test_split, GridSearchCV
 from sklearn.metrics import confusion_matrix, mean_squared_error
-from sklearn.grid_search import GridSearchCV
 from sklearn.datasets import load_iris, load_digits, load_boston
 
 rng = np.random.RandomState(31337)
@@ -22,8 +18,8 @@ print("Zeros and Ones from the Digits dataset: binary classification")
 digits = load_digits(2)
 y = digits['target']
 X = digits['data']
-kf = KFold(y.shape[0], n_folds=2, shuffle=True, random_state=rng)
-for train_index, test_index in kf:
+kf = KFold(n_splits=2, shuffle=True, random_state=rng)
+for train_index, test_index in kf.split(X):
     xgb_model = xgb.XGBClassifier().fit(X[train_index],y[train_index])
     predictions = xgb_model.predict(X[test_index])
     actuals = y[test_index]
@@ -33,8 +29,8 @@ print("Iris: multiclass classification")
 iris = load_iris()
 y = iris['target']
 X = iris['data']
-kf = KFold(y.shape[0], n_folds=2, shuffle=True, random_state=rng)
-for train_index, test_index in kf:
+kf = KFold(n_splits=2, shuffle=True, random_state=rng)
+for train_index, test_index in kf.split(X):
     xgb_model = xgb.XGBClassifier().fit(X[train_index],y[train_index])
     predictions = xgb_model.predict(X[test_index])
     actuals = y[test_index]
@@ -44,8 +40,8 @@ print("Boston Housing: regression")
 boston = load_boston()
 y = boston['target']
 X = boston['data']
-kf = KFold(y.shape[0], n_folds=2, shuffle=True, random_state=rng)
-for train_index, test_index in kf:
+kf = KFold(n_splits=2, shuffle=True, random_state=rng)
+for train_index, test_index in kf.split(X):
     xgb_model = xgb.XGBRegressor().fit(X[train_index],y[train_index])
     predictions = xgb_model.predict(X[test_index])
     actuals = y[test_index]