fix: MaskedCausalVisionTransformer build on all supported timm versions - #2043
Merged
gabrielfruet merged 2 commits intoAug 21, 2026
Conversation
MaskedCausalBlock forwarded its full block-level kwargs to MaskedCausalAttention, including arguments the timm Attention constructor does not accept (e.g. mlp_ratio). This made MaskedCausalVisionTransformer, and therefore the AIM model, fail to construct on every supported timm version (0.9.9-1.0.28) with "Attention.__init__() got an unexpected keyword argument 'mlp_ratio'". Forward only the arguments that Attention.__init__ defines, selected via its signature so the fix stays correct as timm evolves. Add a regression test covering the block and the vision transformer.
timm's VisionTransformer asserts `class_token or global_pool != 'token'`. The AIM examples and benchmark build MaskedCausalVisionTransformer with class_token=False but did not set global_pool, so construction failed on current timm. AIM's self-supervised path uses forward_features and is not affected by global_pool, so "avg" satisfies the assertion without changing behaviour. Regenerate the AIM example notebooks accordingly.
lorinczszabolcs
force-pushed
the
fix-masked-causal-vit-attention-kwargs
branch
from
August 21, 2026 18:29
6ca463a to
9ee3d4e
Compare
Contributor
Author
|
Addressed the failing tests. |
gabrielfruet
approved these changes
Aug 21, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #2042.
Two commits:
Forward only attention kwargs in
MaskedCausalBlock. The block forwarded all of its block-level kwargs to thetimmAttentionlayer.Attentionrejects arguments likemlp_ratio, soMaskedCausalVisionTransformerfailed to build on every supportedtimmversion. This is a regression from MaskedCausalBlock.__init__() got an unexpected keyword argument 'proj_bias' #1841 (the**kwargsforwarding that fixed MaskedCausalBlock.__init__() got an unexpected keyword argument 'proj_bias' #1829). The fix selects the kwargs fromAttention.__init__'s signature, so it keeps MaskedCausalBlock.__init__() got an unexpected keyword argument 'proj_bias' #1841's version-robustness without a hardcoded argument list. Adds a regression test for the block and the model, this backbone had none before.Set
global_pool="avg"in the AIM examples and benchmark. The examples build the backbone withclass_token=Falsebut noglobal_pool, which tripstimm'sclass_token or global_pool != "token"assertion at init, before the block bug is even reached. AIM's self-supervised path usesforward_features, so"avg"satisfies the assertion without changing behavior. Regenerated the notebooks.Verified on
timm0.9.9,1.0.14, and1.0.28.make format-checkpasses and the notebooks are up to date.