From b7088865a4878181703e0cdbe5513792b5046409 Mon Sep 17 00:00:00 2001 From: Sajal Kumar Jana Date: Thu, 5 Nov 2026 05:51:50 +0000 Subject: [PATCH 1/2] fix(scalarization): Register tensor attributes as buffers --- CHANGELOG.md | 7 +++++++ src/torchjd/scalarization/_constant.py | 4 +++- src/torchjd/scalarization/_cosmos.py | 4 +++- src/torchjd/scalarization/_pbi.py | 7 +++++-- src/torchjd/scalarization/_stch.py | 7 +++++-- tests/unit/scalarization/test_constant.py | 8 ++++++++ tests/unit/scalarization/test_cosmos.py | 8 ++++++++ tests/unit/scalarization/test_pbi.py | 16 ++++++++++++++++ tests/unit/scalarization/test_stch.py | 17 +++++++++++++++++ 9 files changed, 72 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ad98060b4..f0775d953 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,13 @@ changelog does not include internal changes that do not affect the user. ## [Unreleased] +### Fixed + +- Fixed `.to()` not moving the tensors given to the `Constant`, `STCH`, `PBI` and `COSMOS` + scalarizers (`weights` and `reference`). These tensors are now registered as non-persistent + buffers, so they follow the dtype and device of the scalarizer without being added to its + `state_dict`. + ## [0.18.1] - 2026-11-02 ### Added diff --git a/src/torchjd/scalarization/_constant.py b/src/torchjd/scalarization/_constant.py index 5223f289d..0d46781f9 100644 --- a/src/torchjd/scalarization/_constant.py +++ b/src/torchjd/scalarization/_constant.py @@ -14,9 +14,11 @@ class Constant(Scalarizer): passed at call time. """ + weights: Tensor + def __init__(self, weights: Tensor) -> None: super().__init__() - self.weights = weights + self.register_buffer("weights", weights, persistent=False) def forward(self, values: Tensor, /) -> Tensor: if values.shape != self.weights.shape: diff --git a/src/torchjd/scalarization/_cosmos.py b/src/torchjd/scalarization/_cosmos.py index 2de5077c7..9135ed7c1 100644 --- a/src/torchjd/scalarization/_cosmos.py +++ b/src/torchjd/scalarization/_cosmos.py @@ -42,6 +42,8 @@ class COSMOS(Scalarizer): objective. """ + weights: Tensor + def __init__(self, lambda_: float, weights: Tensor) -> None: if lambda_ < 0.0: raise ValueError( @@ -50,7 +52,7 @@ def __init__(self, lambda_: float, weights: Tensor) -> None: super().__init__() self.lambda_ = lambda_ - self.weights = weights + self.register_buffer("weights", weights, persistent=False) def forward(self, values: Tensor, /) -> Tensor: check_same_shape_as_values(self.weights, values, "weights") diff --git a/src/torchjd/scalarization/_pbi.py b/src/torchjd/scalarization/_pbi.py index fe8f50255..fb6734a96 100644 --- a/src/torchjd/scalarization/_pbi.py +++ b/src/torchjd/scalarization/_pbi.py @@ -51,14 +51,17 @@ class PBI(Scalarizer): :math:`10^{-6}` at that point and is negligible elsewhere. """ + weights: Tensor + reference: Tensor | None + def __init__(self, theta: float, weights: Tensor, reference: Tensor | None = None) -> None: if theta < 0.0: raise ValueError(f"Parameter `theta` should be non-negative. Found `theta = {theta}`.") super().__init__() self.theta = theta - self.weights = weights - self.reference = reference + self.register_buffer("weights", weights, persistent=False) + self.register_buffer("reference", reference, persistent=False) def forward(self, values: Tensor, /) -> Tensor: check_same_shape_as_values(self.weights, values, "weights") diff --git a/src/torchjd/scalarization/_stch.py b/src/torchjd/scalarization/_stch.py index 25ec95bc1..019dd190b 100644 --- a/src/torchjd/scalarization/_stch.py +++ b/src/torchjd/scalarization/_stch.py @@ -38,6 +38,9 @@ class STCH(Scalarizer): is applied. If provided, it must have the same shape as the values passed at call time. """ + weights: Tensor | None + reference: Tensor | None + def __init__( self, mu: float, @@ -49,8 +52,8 @@ def __init__( super().__init__() self.mu = mu - self.weights = weights - self.reference = reference + self.register_buffer("weights", weights, persistent=False) + self.register_buffer("reference", reference, persistent=False) def forward(self, values: Tensor, /) -> Tensor: check_same_shape_as_values(self.weights, values, "weights") diff --git a/tests/unit/scalarization/test_constant.py b/tests/unit/scalarization/test_constant.py index 126b91c01..1ae6f673e 100644 --- a/tests/unit/scalarization/test_constant.py +++ b/tests/unit/scalarization/test_constant.py @@ -61,3 +61,11 @@ def test_representations() -> None: def test_str_with_non_vector_weights() -> None: assert str(Constant(weights=ones_((3, 4)))) == "Constant(weights of shape (3, 4))" assert str(Constant(weights=ones_(()))) == "Constant(weights of shape ())" + + +def test_to_moves_weights() -> None: + s = Constant(weights=tensor_([1.0, 2.0], dtype=torch.float32)).to(torch.float64) + + assert s.weights.dtype == torch.float64 + assert any(buffer is s.weights for buffer in s.buffers()) + assert len(s.state_dict()) == 0 diff --git a/tests/unit/scalarization/test_cosmos.py b/tests/unit/scalarization/test_cosmos.py index d98369d19..399aa8f39 100644 --- a/tests/unit/scalarization/test_cosmos.py +++ b/tests/unit/scalarization/test_cosmos.py @@ -86,3 +86,11 @@ def test_representations() -> None: s = COSMOS(lambda_=0.5, weights=torch.tensor([0.5, 0.5])) assert repr(s) == "COSMOS(lambda_=0.5, weights=tensor([0.5000, 0.5000]))" assert str(s) == "COSMOS" + + +def test_to_moves_weights() -> None: + s = COSMOS(lambda_=0.5, weights=tensor_([0.5, 0.5], dtype=torch.float32)).to(torch.float64) + + assert s.weights.dtype == torch.float64 + assert any(buffer is s.weights for buffer in s.buffers()) + assert len(s.state_dict()) == 0 diff --git a/tests/unit/scalarization/test_pbi.py b/tests/unit/scalarization/test_pbi.py index f7187d27a..4d3e2d1da 100644 --- a/tests/unit/scalarization/test_pbi.py +++ b/tests/unit/scalarization/test_pbi.py @@ -103,3 +103,19 @@ def test_representations() -> None: s = PBI(theta=5.0, weights=torch.tensor([0.5, 0.5])) assert repr(s) == "PBI(theta=5.0, weights=tensor([0.5000, 0.5000]), reference=None)" assert str(s) == "PBI" + + +def test_to_moves_weights_and_reference() -> None: + s = PBI( + theta=5.0, + weights=tensor_([0.5, 0.5], dtype=torch.float32), + reference=tensor_([0.0, 1.0], dtype=torch.float32), + ).to(torch.float64) + + assert s.reference is not None + assert s.weights.dtype == torch.float64 + assert s.reference.dtype == torch.float64 + buffers = list(s.buffers()) + assert any(buffer is s.weights for buffer in buffers) + assert any(buffer is s.reference for buffer in buffers) + assert len(s.state_dict()) == 0 diff --git a/tests/unit/scalarization/test_stch.py b/tests/unit/scalarization/test_stch.py index 659a6a3bb..2bdf7c374 100644 --- a/tests/unit/scalarization/test_stch.py +++ b/tests/unit/scalarization/test_stch.py @@ -80,3 +80,20 @@ def test_representations() -> None: s = STCH(mu=0.5) assert repr(s) == "STCH(mu=0.5, weights=None, reference=None)" assert str(s) == "STCH" + + +def test_to_moves_weights_and_reference() -> None: + s = STCH( + mu=1.0, + weights=tensor_([0.5, 0.5], dtype=torch.float32), + reference=tensor_([0.0, 1.0], dtype=torch.float32), + ).to(torch.float64) + + assert s.weights is not None + assert s.reference is not None + assert s.weights.dtype == torch.float64 + assert s.reference.dtype == torch.float64 + buffers = list(s.buffers()) + assert any(buffer is s.weights for buffer in buffers) + assert any(buffer is s.reference for buffer in buffers) + assert len(s.state_dict()) == 0 From f0504347599cd4edb883aeb744760e0a73b3ad1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Thu, 5 Nov 2026 09:07:39 +0100 Subject: [PATCH 2/2] fix(scalarization): Make FAMO's min_losses buffer non-persistent Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- src/torchjd/scalarization/_famo.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c0ac57915..14286b4f0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,7 +13,7 @@ changelog does not include internal changes that do not affect the user. - Fixed `.to()` not moving the tensors given to the `Constant`, `STCH`, `PBI` and `COSMOS` scalarizers (`weights` and `reference`). These tensors are now registered as non-persistent buffers, so they follow the dtype and device of the scalarizer without being added to its - `state_dict`. + `state_dict`. `FAMO`'s `min_losses` buffer is also no longer in its `state_dict`. - Fixed `.to()` not moving the tensors given to `Constant`, `ConstantWeighting`, `UPGrad`, `UPGradWeighting`, `DualProj`, `DualProjWeighting`, `AlignedMTL`, `AlignedMTLWeighting`, `ConFIG`, `SDMGradWeighting`, `CRMOGMWeighting` and `GradDrop` (`weights`, `pref_vector`, diff --git a/src/torchjd/scalarization/_famo.py b/src/torchjd/scalarization/_famo.py index 7fca1b413..9d4eb294b 100644 --- a/src/torchjd/scalarization/_famo.py +++ b/src/torchjd/scalarization/_famo.py @@ -128,7 +128,7 @@ def __init__( f"Parameter `min_losses` should have shape {tuple(self._w.shape)} (matching the " f"shape of the logits). Found `min_losses.shape = {tuple(min_losses.shape)}`." ) - self.register_buffer("min_losses", min_losses) + self.register_buffer("min_losses", min_losses, persistent=False) self.lr = lr self.weight_decay = weight_decay