From f25ecb1d2b8d5895f9f498f75b493ad52e9f3589 Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Thu, 22 Oct 2026 00:50:00 +0200
Subject: [PATCH] fix(models): apply model_config__* keys when no model_config
 was given

An estimator constructed with only a trainer_config or
preprocessing_config trains from a default config object, but
set_params() dropped nested model_config__* keys for it, so searches
over them trained the default architecture. The default config is now
adopted as model_config and tuned.

Fixes #613
---
 deeptab/models/base.py                       |  5 +++++
 docs/getting_started/faq.md                  |  2 +-
 tests/models/test_estimator_configuration.py | 17 +++++++++++++++++
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/deeptab/models/base.py b/deeptab/models/base.py
index a65fff6d..d11e4fab 100644
--- a/deeptab/models/base.py
+++ b/deeptab/models/base.py
@@ -474,6 +474,11 @@ def set_params(self, **parameters):
                 elif k == "observability_config":
                     self._apply_observability_config(v)
 
+            if model_config_params and self.model_config is None:
+                # No model_config yet: tune the default config that fit() builds
+                # from, and expose it as model_config so that get_params() and
+                # clone() report the tuned values.
+                self.model_config = self.config
             if model_config_params and self.model_config is not None and hasattr(self.model_config, "set_params"):
                 self.model_config.set_params(**model_config_params)
                 self._config_kwargs = self.model_config.get_params(deep=False)
diff --git a/docs/getting_started/faq.md b/docs/getting_started/faq.md
index 8a463ebf..2f4a1c87 100644
--- a/docs/getting_started/faq.md
+++ b/docs/getting_started/faq.md
@@ -451,7 +451,7 @@ search = GridSearchCV(
 search.fit(X_train, y_train)
 ```
 
-Nested keys such as `model_config__d_model` require the corresponding config object. Keep `n_jobs=1` when sharing one accelerator to avoid concurrent fits competing for its memory.
+Nested keys such as `model_config__d_model` need an estimator constructed with config objects. If only `model_config` is omitted, `model_config__*` keys tune the estimator's default model config, which `get_params()` then reports as `model_config`. Keep `n_jobs=1` when sharing one accelerator to avoid concurrent fits competing for its memory.
 
 ### Can I deploy DeepTab models?
 
diff --git a/tests/models/test_estimator_configuration.py b/tests/models/test_estimator_configuration.py
index fe2f8aad..b052425e 100644
--- a/tests/models/test_estimator_configuration.py
+++ b/tests/models/test_estimator_configuration.py
@@ -240,6 +240,23 @@ def test_set_params_invalid_top_level_key_raises(self):
         with pytest.raises(ValueError, match="abc"):
             model.set_params(abc=123)
 
+    def test_nested_model_config_keys_without_model_config_tune_the_default_config(self):
+        model = MLPRegressor(trainer_config=_FAST_TRAINER).set_params(model_config__layer_sizes=[4])
+        assert model.config.layer_sizes == [4]
+        assert model.get_params()["model_config__layer_sizes"] == [4]
+        assert clone(model).config.layer_sizes == [4]
+
+    def test_grid_search_over_nested_model_config_keys_without_model_config(self, tmp_path):
+        from sklearn.model_selection import GridSearchCV
+
+        estimator = MLPRegressor(trainer_config=_FAST_TRAINER, random_state=0)
+        search = GridSearchCV(estimator, {"model_config__layer_sizes": [[4], [16]]}, cv=2)
+        search.fit(
+            X_reg, y_reg, checkpoint_path=str(tmp_path), accelerator="cpu", logger=False, enable_progress_bar=False
+        )
+        best_width = search.best_params_["model_config__layer_sizes"][0]
+        assert search.best_estimator_._task_model.estimator.layers[0].out_features == best_width
+
 
 _TRAINING_FIELDS = {"lr", "lr_patience", "lr_factor", "weight_decay"}