[BUG] The fitted loss function is not persisted: class_weight models fail to load, plain classifiers reload with MSELoss

[ChrisW09]

Describe the bug

build_save_bundle stores task_model_state_dict — which contains the weighted loss's buffers — but
never stores the loss itself. load() rebuilds TaskModel with a default loss, producing two failures:

1. Hard crash. A classifier trained with class_weight (a documented, first-class fit()
   argument) cannot be reloaded: RuntimeError: Unexpected key(s) in state_dict: "loss_fct._loss.pos_weight" (binary) or "loss_fct._loss.weight" (multiclass). Reproduced.
2. Silent corruption. A plain binary classifier does load — but comes back with
   loss_fct = MSELoss, because the bundle stores num_classes=1 and build_default_task_loss then
   picks the regression loss. Reproduced: type(loaded._task_model.loss_fct).__name__ == 'MSELoss'.
   So loaded.fit(..., rebuild=False) continues training a classifier under mean-squared error.
   loss_fct='focal' likewise reloads as MSELoss.

---

save()/load() round trip raises RuntimeError for any classifier trained with class_weight — the fitted loss is not persisted, but its buffers are in the state dict

Where: deeptab/core/serialization.py (215-341)

build_save_bundle stores task_model_state_dict (which contains the weighted loss's pos_weight/weight buffers) but never stores loss_fct; load() rebuilds TaskModel with the default loss, so load_state_dict hits unexpected keys and raises.

Observed: RuntimeError: Error(s) in loading state_dict for TaskModel:
Unexpected key(s) in state_dict: "loss_fct._loss.pos_weight".
Multiclass is the same with "loss_fct._loss.weight". Related silent-wrong from the same root cause: a plain binary classifier does load, but comes back with loss_fct = MSELoss (bundle stores num_classes=1, so build_default_task_loss picks the regression loss) — so loaded.fit(..., rebuild=False) would continue training a classifier under MSE. loss_fct='focal' likewise loads back as MSELoss.

Expected: A model fitted with class_weight='balanced' (a documented, first-class argument of fit()) should save and reload; the reloaded TaskModel should carry the same loss it was trained with.

Repro

import warnings; warnings.simplefilter('ignore')
import numpy as np, pandas as pd
from deeptab.models.mlp import MLPClassifier
from deeptab.configs import MLPConfig, TrainerConfig
rng = np.random.default_rng(0)
X = pd.DataFrame({'a': rng.normal(size=64), 'b': rng.normal(size=64)})
y = np.array([0]*50 + [1]*14)          # imbalanced -> class_weight='balanced'
m = MLPClassifier(model_config=MLPConfig(layer_sizes=[8]),
                  trainer_config=TrainerConfig(max_epochs=1, batch_size=16, patience=2), random_state=0)
m.fit(X, y, class_weight='balanced', accelerator='cpu', devices=1)
m.save('x.deeptab')
MLPClassifier.load('x.deeptab')        # RuntimeError

---

Expected behavior
The save bundle should record enough to rebuild the loss the model was actually trained with (loss
name/spec plus class weights), and the reloaded TaskModel should carry it.

Screenshots
n/a

Desktop (please complete the following information):

• OS: macOS (Darwin 25.5.0, arm64)
• Python version: 3.11.15
• deeptab Version: 2.0.0 (main @ 4e6a359)

Additional context
torch 2.9.1, lightning 2.6.5, scikit-learn 1.9.0, numpy 2.4.6. Found in a second-pass review of v2.0.0
(seven independent lenses, each finding adversarially re-verified by a second reviewer, then re-run by
hand). Distinct from the already-filed #409-#426.