Summary:
load_forecaster_module_from_checkpoint in neural_lam/train_model.py (line 41) is defined but nothing ever calls it. A repo-wide search for the name returns a single hit, which is the definition itself.
The --load / --eval path in main() doesn't use it. Instead main() rebuilds the predictor, forecaster and ForecasterModule inline from the CLI args and then leans on Lightning's ckpt_path=args.load to restore the weights:
- trainer.test(model=model, datamodule=data_module, ckpt_path=args.load) for eval
- trainer.fit(model=model, datamodule=data_module, ckpt_path=args.load) for training
So the helper is genuinely orphaned.
Why this is worth cleaning up
There are two reasonable directions, and I think it's worth a quick decision either way:
- Deleting it. It's dead, and it has already drifted out of sync with main() .
- Wire it in. The function's stated goal is actually a nice usability win: let eval rebuild the architecture straight from the checkpoint's saved hyperparameters instead of forcing the user to re-pass matching --model / --graph / --hidden_dim / ... flags. Right now if those eval flags don't match how the model was trained, you either crash or silently load into the wrong architecture. If we go this way, the helper needs to be fixed first and then routed into the --eval --load path.
The helper is also stale, not just unused
If we do keep it, note that it's already behind main(). The inline construction in main() (lines 443-459) passes four GNN-type kwargs:
g2m_gnn_type, m2g_gnn_type, mesh_up_gnn_type, mesh_down_gnn_type
The helper (lines 53-65) stops at output_clamping_upper and never forwards any of them. So for any checkpoint trained with a non-default GNN type, the helper would rebuild the wrong architecture and silently fall back to the defaults. Its docstring promises to recover architecture kwargs automatically, but the implementation no longer does that fully.
Relation to existing issues
Suggested resolution
Unless we want the reconstruct-from-checkpoint feature now, just delete the function. If we do want it, fix the kwarg drift first (ideally by filtering kwargs against the predictor's signature so flat and hierarchical models are both handled, which also fixes the #686 crash) and then call it from the --eval --load branch of main().
Happy to open a PR for whichever direction the maintainers prefer. Assign me the issue if found suitable !
Summary:
load_forecaster_module_from_checkpoint in neural_lam/train_model.py (line 41) is defined but nothing ever calls it. A repo-wide search for the name returns a single hit, which is the definition itself.
The --load / --eval path in main() doesn't use it. Instead main() rebuilds the predictor, forecaster and ForecasterModule inline from the CLI args and then leans on Lightning's ckpt_path=args.load to restore the weights:
So the helper is genuinely orphaned.
Why this is worth cleaning up
There are two reasonable directions, and I think it's worth a quick decision either way:
The helper is also stale, not just unused
If we do keep it, note that it's already behind main(). The inline construction in main() (lines 443-459) passes four GNN-type kwargs:
g2m_gnn_type, m2g_gnn_type, mesh_up_gnn_type, mesh_down_gnn_type
The helper (lines 53-65) stops at output_clamping_upper and never forwards any of them. So for any checkpoint trained with a non-default GNN type, the helper would rebuild the wrong architecture and silently fall back to the defaults. Its docstring promises to recover architecture kwargs automatically, but the implementation no longer does that fully.
Relation to existing issues
Suggested resolution
Unless we want the reconstruct-from-checkpoint feature now, just delete the function. If we do want it, fix the kwarg drift first (ideally by filtering kwargs against the predictor's signature so flat and hierarchical models are both handled, which also fixes the #686 crash) and then call it from the --eval --load branch of main().
Happy to open a PR for whichever direction the maintainers prefer. Assign me the issue if found suitable !