BaseLoss learnable variance parameters are not checkpointed

[djgagne]

Summary

BaseLoss with var_weighting: "learnable" (Kendall–Gal uncertainty weighting) owns a trainable per-variable log σ² parameter, BaseLoss.log_variance. train_gen2 adds it to the optimizer, so training works — but the parameter is not saved to or restored from checkpoints. It re-initializes from the scaler statistics on every resume.

Introduced with the Gen 2 loss framework in #472, where it is disclosed in the PR description and in the learnable sections of credit/losses/base.py and docs/source/Losses.md.

Impact

Resuming a learnable run silently discards the learned weighting and restarts it from the scaler-derived initialization. The loss value jumps at the resume boundary and the effective per-variable weighting regresses. Nothing errors, so it is easy to miss.

Only affects var_weighting: "learnable". The static modes (inverse_variance, manual, none) hold no state and are unaffected.

What needs doing

• Save criterion.log_variance (and any future BaseLoss state) in BaseTrainer.save_checkpoint, and restore it on resume.
• Decide the FSDP/FSDP2 story — the parameter is tiny and replicated, so it likely wants to be gathered on rank 0 and broadcast on load rather than sharded.
• The validation criterion is a separate BaseLoss instance; confirm whether it needs the same state or should stay at init values.
• A test that trains a step, saves, reloads, and asserts log_variance round-trips.

Workaround

Use a static var_weighting for runs that will be resumed, or accept the re-initialization if resuming rarely.