diff --git a/sklearn/model_selection/_search.py b/sklearn/model_selection/_search.py index 424263c0d1c3d..0d0c15c54c468 100644 --- a/sklearn/model_selection/_search.py +++ b/sklearn/model_selection/_search.py @@ -603,7 +603,8 @@ def _store(key_name, array, weights=None, splits=False, rank=False): _store('test_score', test_scores, splits=True, rank=True, weights=test_sample_counts if self.iid else None) - _store('train_score', train_scores, splits=True) + if self.return_train_score: + _store('train_score', train_scores, splits=True) _store('fit_time', fit_time) _store('score_time', score_time) diff --git a/sklearn/model_selection/tests/test_search.py b/sklearn/model_selection/tests/test_search.py index fa4949d317052..30daacf4c42fb 100644 --- a/sklearn/model_selection/tests/test_search.py +++ b/sklearn/model_selection/tests/test_search.py @@ -1140,3 +1140,13 @@ def test_stochastic_gradient_loss_param(): assert_false(hasattr(clf, "predict_proba")) clf.fit(X, y) assert_false(hasattr(clf, "predict_proba")) + + +def test_search_train_scores_set_to_false(): + X = np.arange(6).reshape(6, -1) + y = [0, 0, 0, 1, 1, 1] + clf = LinearSVC(random_state=0) + + gs = GridSearchCV(clf, param_grid={'C': [0.1, 0.2]}, + return_train_score=False) + gs.fit(X, y)