From 910c2f2a5ac260807694e6cae302ac277eb5dcef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Wed, 7 Oct 2026 18:26:18 +0200 Subject: [PATCH] Avoid tensor creation in autogram hook The forward hook of each module built a CPU tensor of output sizes and called argmin().item() on it to find its smallest output. It now uses Python's min with a key, which selects the same output (the first one in case of ties) without creating a tensor at every forward pass. Co-Authored-By: Claude Opus 5.5 --- src/torchjd/autogram/_module_hook_manager.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/torchjd/autogram/_module_hook_manager.py b/src/torchjd/autogram/_module_hook_manager.py index d04372a92..1ea45ae84 100644 --- a/src/torchjd/autogram/_module_hook_manager.py +++ b/src/torchjd/autogram/_module_hook_manager.py @@ -1,5 +1,5 @@ import weakref -from typing import Any, cast +from typing import Any import torch from torch import Tensor, nn @@ -129,9 +129,8 @@ def __call__( # We only care about running the AutogramNode, so we need one of its child # edges (the edges of the original outputs of the model) as target. For memory # efficiency, we select the smallest one (that requires grad). - preference = torch.tensor([t.numel() for t in rg_outputs]) - index = cast(int, preference.argmin().item()) - self.target_edges.register(get_gradient_edge(rg_outputs[index])) + smallest_rg_output = min(rg_outputs, key=lambda t: t.numel()) + self.target_edges.register(get_gradient_edge(smallest_rg_output)) autograd_fn_rg_outputs = AutogramNode.apply( self.gramian_accumulation_phase,