Fix mypy warnings for timm 1.0.16

Signed-off-by: Beat Buesser <[email protected]>

Author: beat-buesser

art/estimators/certification/derandomized_smoothing/pytorch.py

@@ -406,12 +406,15 @@ def get_models(cls, generate_from_null: bool = False) -> list[str]:
         return supported
 
     @staticmethod
-    def create_vision_transformer(variant: str, pretrained: bool = False, **kwargs) -> "PyTorchVisionTransformer":
+    def create_vision_transformer(
+        variant: str, pretrained: bool = False, use_naflex: bool | None = None, **kwargs
+    ) -> "PyTorchVisionTransformer":
         """
         Creates a vision transformer using PyTorchViT which controls the forward pass of the model
 
         :param variant: The name of the vision transformer to load
         :param pretrained: If to load pre-trained weights
+        :param use_naflex:  If using NaFlexVit.
         :return: A ViT with the required methods needed for ART
         """
 

art/estimators/certification/derandomized_smoothing/vision_transformers/pytorch.py

@@ -167,11 +167,12 @@ def drop_tokens(x: torch.Tensor, indexes: torch.Tensor) -> torch.Tensor:
         x_no_cl = torch.reshape(x_no_cl, shape=(shape[0], -1, shape[-1]))
         return torch.cat((cls_token, x_no_cl), dim=1)
 
-    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+    def forward_features(self, x: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor:
         """
         The forward pass of the ViT.
 
         :param x: Input data.
+        :param attn_mask: Attention mask.
         :return: The input processed by the ViT backbone
         """