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 = (