From 918ccbf36df7433e5030df465b0bd61f25b06027 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Mon, 29 Sep 2025 13:47:46 +0200 Subject: [PATCH 1/2] doc(autogram): Add warning about using child parameters --- src/torchjd/autogram/_engine.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/torchjd/autogram/_engine.py b/src/torchjd/autogram/_engine.py index 6a3fbe03..7c1a9f46 100644 --- a/src/torchjd/autogram/_engine.py +++ b/src/torchjd/autogram/_engine.py @@ -147,6 +147,20 @@ class Engine: The alternative is to use ``batch_dim=None``, but it's not recommended since it will increase memory usage by a lot and thus typically slow down computation. + .. warning:: + Parent modules should call their child modules directly rather than using their child + modules' parameters themselves. For instance, the following model is not be supported: + + >>> class Model(nn.Module): + >>> def __init__(self): + >>> super().__init__() + >>> self.linear = nn.Linear(2, 3) # Child module + >>> + >>> def forward(self, input: Tensor) -> Tensor: + >>> # Incorrect: Use the child module's parameters directly without calling it. + >>> return input @ self.linear.weight.T + self.linear.bias + >>> # Correct alternative: return self.linear(input) + .. note:: For maximum efficiency, modules should ideally not contain both direct trainable parameters and child modules, especially if those direct trainable parameters are used From 0ad52440338fe0dac9937c92063d65b96ba9da88 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= <31951177+ValerianRey@users.noreply.github.com> Date: Mon, 29 Sep 2025 15:00:46 +0200 Subject: [PATCH 2/2] Update src/torchjd/autogram/_engine.py Co-authored-by: Pierre Quinton --- src/torchjd/autogram/_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/torchjd/autogram/_engine.py b/src/torchjd/autogram/_engine.py index 7c1a9f46..be7104f0 100644 --- a/src/torchjd/autogram/_engine.py +++ b/src/torchjd/autogram/_engine.py @@ -149,7 +149,7 @@ class Engine: .. warning:: Parent modules should call their child modules directly rather than using their child - modules' parameters themselves. For instance, the following model is not be supported: + modules' parameters themselves. For instance, the following model is not supported: >>> class Model(nn.Module): >>> def __init__(self):