Fix mse in parameter search (#313)

NicolasHug · Nov 2, 2019 · 74e4f06 · 74e4f06
1 parent d29b255
commit 74e4f06
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/surprise/model_selection/search.py b/surprise/model_selection/search.py
@@ -148,7 +148,7 @@ def fit(self, data):
                 len(indices)) + 1  # sklearn starts rankings at 1 as well.
 
             # set best_index, and best_xxxx attributes
-            if m in ('mae', 'rmse'):
+            if m in ('mae', 'rmse', 'mse'):
                 best_index[m] = mean_test_measures.argmin()
             elif m in ('fcp',):
                 best_index[m] = mean_test_measures.argmax()