Use ModelSpec.cross_validate() in ModelSelectionGenerationNode

Summary: See title

Reviewed By: ldworkin

Differential Revision: D32838687

fbshipit-source-id: 6ac6829d75d514226aa5080fb4f44ed2a2b6a639

ax/modelbridge/generation_node.py

@@ -17,11 +17,7 @@
 from ax.core.search_space import SearchSpace
 from ax.exceptions.core import UserInputError
 from ax.modelbridge.base import ModelBridge
-from ax.modelbridge.cross_validation import (
-    cross_validate,
-    compute_diagnostics,
-    BestModelSelector,
-)
+from ax.modelbridge.cross_validation import BestModelSelector
 from ax.modelbridge.model_spec import ModelSpec, FactoryFunctionModelSpec
 from ax.modelbridge.registry import (
     ModelRegistryBase,
@@ -50,7 +46,6 @@ def __init__(
         self,
         model_specs: List[ModelSpec],
         best_model_selector: Optional[BestModelSelector] = None,
-        cvkwargs: Optional[Dict[str, Any]] = None,
     ) -> None:
         # While `GenerationNode` only handles a single `ModelSpec` in the `gen`
         # and `_pick_fitted_model_to_gen_from` methods, we validate the
@@ -59,7 +54,6 @@ def __init__(
         # method to bypass that validation.
         self.model_specs = model_specs
         self.best_model_selector = best_model_selector
-        self.cvkwargs = cvkwargs
 
     def fit(
         self,
@@ -135,14 +129,11 @@ def _pick_fitted_model_to_gen_from(self) -> ModelSpec:
                 raise NotImplementedError(CANNOT_SELECT_ONE_MODEL_MSG)
             return self.model_specs[0]
 
-        cvkwargs = self.cvkwargs or {}
-        cv_diagnostics = []
         for model_spec in self.model_specs:
-            cv_result = cross_validate(model_spec.fitted_model, **cvkwargs)
-            cv_diagnostics.append(compute_diagnostics(cv_result))
+            model_spec.cross_validate()
 
         best_model_index = not_none(self.best_model_selector).best_diagnostic(
-            cv_diagnostics
+            diagnostics=[not_none(m.diagnostics) for m in self.model_specs],
         )
         return self.model_specs[best_model_index]