Skip to content

Commit 6d0b3a8

Browse files
authored
refactor(autogram): Fix outdated list of incompatible modules (#450)
* 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.
1 parent a211197 commit 6d0b3a8

File tree

1 file changed

+0
-5
lines changed

1 file changed

+0
-5
lines changed

src/torchjd/autogram/_engine.py

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,6 @@
1919
nn.LazyBatchNorm3d,
2020
nn.SyncBatchNorm,
2121
nn.RNNBase,
22-
nn.Transformer,
23-
nn.TransformerEncoder,
24-
nn.TransformerDecoder,
25-
nn.TransformerEncoderLayer,
26-
nn.TransformerDecoderLayer,
2722
)
2823

2924
_TRACK_RUNNING_STATS_MODULE_TYPES = (

0 commit comments

Comments
 (0)