From c439f52b8856abd49427e293ee91b30c02641b2e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Fri, 10 Oct 2025 17:05:20 +0200 Subject: [PATCH] refactor(autogram): Fix outdated list of incompatible modules * Transformer and helper modules are now compatible with batched engine. It somehow worked even with those modules in the list because Transformer and helper modules do not have direct params and thus do not get hooked themselves anyway. --- src/torchjd/autogram/_engine.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/torchjd/autogram/_engine.py b/src/torchjd/autogram/_engine.py index a60bd4c4..9cab68f1 100644 --- a/src/torchjd/autogram/_engine.py +++ b/src/torchjd/autogram/_engine.py @@ -19,11 +19,6 @@ nn.LazyBatchNorm3d, nn.SyncBatchNorm, nn.RNNBase, - nn.Transformer, - nn.TransformerEncoder, - nn.TransformerDecoder, - nn.TransformerEncoderLayer, - nn.TransformerDecoderLayer, ) _TRACK_RUNNING_STATS_MODULE_TYPES = (