diff --git a/src/torchjd/aggregation/_cagrad.py b/src/torchjd/aggregation/_cagrad.py index b2da3efe..40a5d5b0 100644 --- a/src/torchjd/aggregation/_cagrad.py +++ b/src/torchjd/aggregation/_cagrad.py @@ -19,8 +19,6 @@ # Non-differentiable: the cvxpy solver operates on numpy arrays, breaking the autograd graph. class CAGradWeighting(_WithOptionalDeps, _GramianWeighting, _NonDifferentiable): - _REQUIRED_DEPS = ["numpy", "cvxpy", "clarabel"] - _INSTALL_HINT = 'Install them with: pip install "torchjd[cagrad]"' """ :class:`~torchjd.aggregation.Weighting` [:class:`~torchjd.linalg.PSDMatrix`] giving the weights of :class:`~torchjd.aggregation.CAGrad`. @@ -37,6 +35,9 @@ class CAGradWeighting(_WithOptionalDeps, _GramianWeighting, _NonDifferentiable): function. """ + _REQUIRED_DEPS = ["numpy", "cvxpy", "clarabel"] + _INSTALL_HINT = 'Install them with: pip install "torchjd[cagrad]"' + def __init__(self, c: float, norm_eps: float = 0.0001) -> None: super().__init__() self.c = c diff --git a/src/torchjd/aggregation/_nash_mtl.py b/src/torchjd/aggregation/_nash_mtl.py index d3f11a9c..a4583e6c 100644 --- a/src/torchjd/aggregation/_nash_mtl.py +++ b/src/torchjd/aggregation/_nash_mtl.py @@ -22,8 +22,6 @@ # Non-differentiable: the cvxpy solver operates on numpy arrays, breaking the autograd graph. class _NashMTLWeighting(_WithOptionalDeps, _MatrixWeighting, Stateful, _NonDifferentiable): - _REQUIRED_DEPS = ["numpy", "cvxpy", "ecos"] - _INSTALL_HINT = 'Install them with: pip install "torchjd[nash_mtl]"' """ :class:`~torchjd.Stateful` :class:`~torchjd.aggregation.Weighting` [:class:`~torchjd.linalg.Matrix`] that @@ -45,6 +43,9 @@ class _NashMTLWeighting(_WithOptionalDeps, _MatrixWeighting, Stateful, _NonDiffe affects the shape of the cached state). """ + _REQUIRED_DEPS = ["numpy", "cvxpy", "ecos"] + _INSTALL_HINT = 'Install them with: pip install "torchjd[nash_mtl]"' + def __init__( self, n_tasks: int,