From 3619bd386b66b5131c5b78e21e6775982d2b984d Mon Sep 17 00:00:00 2001
From: Raghav RV <rvraghav93@gmail.com>
Date: Fri, 30 Sep 2016 05:53:25 +0200
Subject: [PATCH] BUGFIX _store train_scores only if return_train_score is True
(#7535)
---
sklearn/model_selection/_search.py | 3 ++-
sklearn/model_selection/tests/test_search.py | 10 ++++++++++
2 files changed, 12 insertions(+), 1 deletion(-)
diff --git a/sklearn/model_selection/_search.py b/sklearn/model_selection/_search.py
index 424263c0d1..0d0c15c54c 100644
--- a/sklearn/model_selection/_search.py
+++ b/sklearn/model_selection/_search.py
@@ -603,7 +603,8 @@ class BaseSearchCV(six.with_metaclass(ABCMeta, BaseEstimator,
_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 fa4949d317..30daacf4c4 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)
--
GitLab