[SajalDevX]

Follow-up to #826, as suggested there by @ValerianRey.

Tensors given to aggregators and weightings were stored as plain attributes, so .to() did not move them. For example Constant(torch.tensor([1., 2.])).to(torch.float64) kept its weights in float32.

This registers every such tensor as a buffer instead:

• ConstantWeighting.weights (and Constant now reads its weights from its weighting instead of keeping a second reference in _weights)
• _pref_vector of UPGradWeighting, DualProjWeighting, AlignedMTLWeighting, ConFIG and SDMGradWeighting (the aggregators wrapping them get it for free)
• _initial_weights of CRMOGMWeighting

Attribute access is unchanged (weights, pref_vector, the setters, repr/str). The stateful state tensors (_w, _lambda, _phi_t, ...) were already buffers, so I left them alone.

Open question: persistent=False. I went with non persistent buffers. These tensors are constructor arguments, like lr or alpha which are not in the state_dict either, and making them persistent would add new keys to the state_dict of any module holding an aggregator, so loading a checkpoint saved with 0.18 using strict=True would fail with missing keys. The downside is that a checkpoint does not restore them, which is a bit different from the stateful buffers (_w, _lambda, ...) that are persistent. Happy to switch to persistent=True if you prefer that.

Not included: the scalarizers in torchjd.scalarization (Constant, STCH, PBI, COSMOS) also keep weights / reference as plain attributes. I can do the same there in this PR or a separate one if you want.

Testing

Added a test_to_moves_* test for each class above. Each builds the object with float32 tensors, calls .to(torch.float64), and checks that the tensors are now float64, that they are among the module's buffers, and that they are not in the state_dict. The 7 new tests fail on main and pass with this change.

• uv run pytest tests/unit -W error: 3703 passed, 66 skipped, 33 xfailed (CPU)
• new tests also pass with PYTEST_TORCH_DTYPE=float64
• ruff check, ruff format, pre-commit on changed files: clean
• ty check: no new diagnostics (only the existing unresolved imports for the optional plot deps I did not install)
• not run: CUDA tests (no GPU here)

[ValerianRey]

Hi, thanks for the PR!

Please correct me if I'm wrong, but I think GradDrop's leak param is missing. Apart from that, LGTM.

Also, could you please open a PR with a similar change for the scalarizers?