birder 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

birder/adversarial/deepfool.py

@@ -2,6 +2,8 @@
 DeepFool
 
 Paper "DeepFool: a simple and accurate method to fool deep neural networks", https://arxiv.org/abs/1511.04599
+
+Generated by gpt-5.2-codex xhigh.
 """
 
 from typing import Optional

birder/adversarial/simba.py

@@ -2,6 +2,8 @@
 SimBA (Simple Black-box Attack)
 
 Paper "Simple Black-box Adversarial Attacks", https://arxiv.org/abs/1905.07121
+
+Generated by gpt-5.2-codex xhigh.
 """
 
 from typing import Optional

birder/common/masking.py

@@ -84,8 +84,8 @@ def mask_tensor(
 
     (B, H, W, _) = x.size()
 
-    shaped_mask = mask.reshape(-1, H // patch_factor, W // patch_factor)
-    shaped_mask = shaped_mask.repeat_interleave(patch_factor, axis=1).repeat_interleave(patch_factor, axis=2)
+    shaped_mask = mask.reshape(B, H // patch_factor, W // patch_factor)
+    shaped_mask = shaped_mask.repeat_interleave(patch_factor, dim=1).repeat_interleave(patch_factor, dim=2)
     shaped_mask = shaped_mask.unsqueeze(3).type_as(x)
 
     if mask_token is not None:
@@ -228,14 +228,23 @@ class Masking:
 
 
 class UniformMasking(Masking):
-    def __init__(self, input_size: tuple[int, int], mask_ratio: float, device: Optional[torch.device] = None) -> None:
+    def __init__(
+        self,
+        input_size: tuple[int, int],
+        mask_ratio: float,
+        min_mask_size: int = 1,
+        device: Optional[torch.device] = None,
+    ) -> None:
         self.h = input_size[0]
         self.w = input_size[1]
         self.mask_ratio = mask_ratio
+        self.min_mask_size = min_mask_size
         self.device = device
 
     def __call__(self, batch_size: int) -> torch.Tensor:
-        return uniform_mask(batch_size, self.h, self.w, self.mask_ratio, device=self.device)[0]
+        return uniform_mask(
+            batch_size, self.h, self.w, self.mask_ratio, min_mask_size=self.min_mask_size, device=self.device
+        )[0]
 
 
 class BlockMasking(Masking):

birder/inference/classification.py

@@ -85,7 +85,7 @@ def infer_batch(
             logits = net(t(tta_input), **kwargs)
             outs.append(logits if return_logits is True else F.softmax(logits, dim=1))
 
-        out = torch.stack(outs).mean(axis=0)
+        out = torch.stack(outs).mean(dim=0)
 
     else:
         logits = net(inputs, **kwargs)

birder/introspection/__init__.py

@@ -1,5 +1,6 @@
 from birder.introspection.attention_rollout import AttentionRollout
 from birder.introspection.base import InterpretabilityResult
+from birder.introspection.feature_pca import FeaturePCA
 from birder.introspection.gradcam import GradCAM
 from birder.introspection.guided_backprop import GuidedBackprop
 from birder.introspection.transformer_attribution import TransformerAttribution
@@ -7,6 +8,7 @@ from birder.introspection.transformer_attribution import TransformerAttribution
 __all__ = [
     "InterpretabilityResult",
     "AttentionRollout",
+    "FeaturePCA",
     "GradCAM",
     "GuidedBackprop",
     "TransformerAttribution",

birder/introspection/base.py

@@ -2,7 +2,6 @@ from collections.abc import Callable
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
-from typing import Protocol
 
 import matplotlib
 import matplotlib.pyplot as plt
@@ -27,12 +26,6 @@ class InterpretabilityResult:
         plt.show()
 
 
-class Interpreter(Protocol):
-    def __call__(
-        self, image: str | Path | Image.Image, target_class: Optional[int] = None
-    ) -> InterpretabilityResult: ...
-
-
 def load_image(image: str | Path | Image.Image) -> Image.Image:
     if isinstance(image, (str, Path)):
         return Image.open(image)

birder/introspection/feature_pca.py

@@ -0,0 +1,101 @@
+from collections.abc import Callable
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+from PIL import Image
+from sklearn.decomposition import PCA
+
+from birder.introspection.base import InterpretabilityResult
+from birder.introspection.base import preprocess_image
+from birder.net.base import DetectorBackbone
+
+
+class FeaturePCA:
+    """
+    Visualizes feature maps using Principal Component Analysis
+
+    This method extracts feature maps from a specified stage of a DetectorBackbone model,
+    applies PCA to reduce the channel dimension to 3 components, and visualizes them as an RGB image where:
+    - R channel = 1st principal component (most important)
+    - G channel = 2nd principal component
+    - B channel = 3rd principal component
+    """
+
+    def __init__(
+        self,
+        net: DetectorBackbone,
+        device: torch.device,
+        transform: Callable[..., torch.Tensor],
+        normalize: bool = False,
+        channels_last: bool = False,
+        stage: Optional[str] = None,
+    ) -> None:
+        self.net = net.eval()
+        self.device = device
+        self.transform = transform
+        self.normalize = normalize
+        self.channels_last = channels_last
+        self.stage = stage
+
+    def __call__(self, image: str | Path | Image.Image) -> InterpretabilityResult:
+        (input_tensor, rgb_img) = preprocess_image(image, self.transform, self.device)
+
+        with torch.inference_mode():
+            features_dict = self.net.detection_features(input_tensor)
+
+        if self.stage is not None:
+            features = features_dict[self.stage]
+        else:
+            features = list(features_dict.values())[-1]  # Use the last stage by default
+
+        features_np = features.cpu().numpy()
+
+        # Handle channels_last format (B, H, W, C) vs channels_first (B, C, H, W)
+        if self.channels_last is True:
+            (B, H, W, C) = features_np.shape
+            # Already in (B, H, W, C), just reshape to (B*H*W, C)
+            features_reshaped = features_np.reshape(-1, C)
+        else:
+            (B, C, H, W) = features_np.shape
+            # Reshape to (spatial_points, channels) for PCA
+            features_reshaped = features_np.reshape(B, C, -1)
+            features_reshaped = features_reshaped.transpose(0, 2, 1)  # (B, H*W, C)
+            features_reshaped = features_reshaped.reshape(-1, C)  # (B*H*W, C)
+
+        x = features_reshaped
+        if self.normalize is True:
+            x = x / (np.linalg.norm(x, axis=1, keepdims=True) + 1e-6)
+
+        pca = PCA(n_components=3)
+        pca_features = pca.fit_transform(x)
+        pca_features = pca_features.reshape(B, H, W, 3)
+
+        # Extract all 3 components (B=1)
+        pca_rgb = pca_features[0]  # (H, W, 3)
+
+        # Normalize each channel independently to [0, 1]
+        for i in range(3):
+            channel = pca_rgb[:, :, i]
+            channel = channel - channel.min()
+            channel = channel / (channel.max() + 1e-7)
+            pca_rgb[:, :, i] = channel
+
+        target_size = (input_tensor.size(-1), input_tensor.size(-2))  # PIL expects (width, height)
+        pca_rgb_resized = (
+            np.array(
+                Image.fromarray((pca_rgb * 255).astype(np.uint8)).resize(target_size, Image.Resampling.BILINEAR)
+            ).astype(np.float32)
+            / 255.0
+        )
+
+        visualization = (pca_rgb_resized * 255).astype(np.uint8)
+
+        return InterpretabilityResult(
+            original_image=rgb_img,
+            visualization=visualization,
+            raw_output=pca_rgb.astype(np.float32),
+            logits=None,
+            predicted_class=None,
+        )

birder/kernels/soft_nms/soft_nms.cpp

@@ -4,6 +4,9 @@
 * Taken from:
 * https://github.com/MrParosk/soft_nms
 * Licensed under the MIT License
+*
+* Modified by:
+* Ofer Hasson — 2026-01-10
 **************************************************************************************************
 */
 
@@ -40,8 +43,8 @@ torch::Tensor calculate_iou(const torch::Tensor& boxes, const torch::Tensor& are
     auto xx2 = torch::minimum(boxes.index({idx, 2}), boxes.index({Slice(idx + 1, None), 2}));
     auto yy2 = torch::minimum(boxes.index({idx, 3}), boxes.index({Slice(idx + 1, None), 3}));
 
-    auto w = torch::maximum(torch::zeros_like(xx1), xx2 - xx1);
-    auto h = torch::maximum(torch::zeros_like(yy1), yy2 - yy1);
+    auto w = (xx2 - xx1).clamp_min(0);
+    auto h = (yy2 - yy1).clamp_min(0);
 
     auto intersection = w * h;
     auto union_ = areas.index({idx}) + areas.index({Slice(idx + 1, None)}) - intersection;

birder/model_registry/model_registry.py

@@ -87,14 +87,15 @@ class ModelRegistry:
         no further registration is needed.
         """
 
+        alias_key = alias.lower()
         if net_type.auto_register is False:
             # Register the model manually, as the base class doesn't take care of that for us
-            registry.register_model(alias, type(alias, (net_type,), {"config": config}))
+            self.register_model(alias_key, type(alias, (net_type,), {"config": config}))
 
         if alias in self.aliases:
             warnings.warn(f"Alias {alias} is already registered", UserWarning)
 
-        self.aliases[alias] = type(alias, (net_type,), {"config": config})
+        self.aliases[alias_key] = type(alias, (net_type,), {"config": config})
 
     def register_weights(self, name: str, weights_info: manifest.ModelMetadataType) -> None:
         if name in self._pretrained_nets:

birder/net/convnext_v1.py

@@ -195,6 +195,21 @@ class ConvNeXt_v1(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
         return self.features(x)
 
 
+registry.register_model_config(
+    "convnext_v1_atto",  # Not in the original v1, taken from v2
+    ConvNeXt_v1,
+    config={"in_channels": [40, 80, 160, 320], "num_layers": [2, 2, 6, 2], "drop_path_rate": 0.0},
+)
+registry.register_model_config(
+    "convnext_v1_femto",  # Not in the original v1, taken from v2
+    ConvNeXt_v1,
+    config={"in_channels": [48, 96, 192, 384], "num_layers": [2, 2, 6, 2], "drop_path_rate": 0.0},
+)
+registry.register_model_config(
+    "convnext_v1_pico",  # Not in the original v1, taken from v2
+    ConvNeXt_v1,
+    config={"in_channels": [64, 128, 256, 512], "num_layers": [2, 2, 6, 2], "drop_path_rate": 0.0},
+)
 registry.register_model_config(
     "convnext_v1_nano",  # Not in the original v1, taken from v2
     ConvNeXt_v1,
@@ -220,6 +235,11 @@ registry.register_model_config(
     ConvNeXt_v1,
     config={"in_channels": [192, 384, 768, 1536], "num_layers": [3, 3, 27, 3], "drop_path_rate": 0.5},
 )
+registry.register_model_config(
+    "convnext_v1_huge",  # Not in the original v1, taken from v2
+    ConvNeXt_v1,
+    config={"in_channels": [352, 704, 1408, 2816], "num_layers": [3, 3, 27, 3], "drop_path_rate": 0.5},
+)
 
 registry.register_weights(
     "convnext_v1_tiny_eu-common256px",

birder/net/fastvit.py

@@ -607,7 +607,6 @@ class AttentionBlock(nn.Module):
 
 
 class FastVitStage(nn.Module):
-    # pylint: disable=too-many-arguments,too-many-positional-arguments
     def __init__(
         self,
         dim: int,

birder/net/flexivit.py

@@ -98,6 +98,8 @@ class FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         layer_scale_init_value: Optional[float] = self.config.get("layer_scale_init_value", None)
         pre_norm: bool = self.config.get("pre_norm", False)
         post_norm: bool = self.config.get("post_norm", True)
+        qkv_bias: bool = self.config.get("qkv_bias", True)
+        qk_norm: bool = self.config.get("qk_norm", False)
         num_reg_tokens: int = self.config.get("num_reg_tokens", 0)
         class_token: bool = self.config.get("class_token", True)
         attn_pool_head: bool = self.config.get("attn_pool_head", False)
@@ -186,6 +188,8 @@ class FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
             attention_dropout,
             dpr,
             pre_norm=pre_norm,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
             activation_layer=act_layer,
             layer_scale_init_value=layer_scale_init_value,
             norm_layer=norm_layer,
@@ -224,6 +228,7 @@ class FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
             drop_path=0,
             activation_layer=act_layer,
             norm_layer=norm_layer,
+            norm_layer_eps=norm_layer_eps,
             mlp_layer=mlp_layer,
         )
 

birder/net/focalnet.py

@@ -245,7 +245,6 @@ class FocalNetBlock(nn.Module):
 
 
 class FocalNetStage(nn.Module):
-    # pylint: disable=too-many-arguments,too-many-positional-arguments
     def __init__(
         self,
         dim: int,

birder/net/rope_flexivit.py

@@ -69,6 +69,8 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         layer_scale_init_value: Optional[float] = self.config.get("layer_scale_init_value", None)
         pre_norm: bool = self.config.get("pre_norm", False)
         post_norm: bool = self.config.get("post_norm", True)
+        qkv_bias: bool = self.config.get("qkv_bias", True)
+        qk_norm: bool = self.config.get("qk_norm", False)
         num_reg_tokens: int = self.config.get("num_reg_tokens", 0)
         class_token: bool = self.config.get("class_token", True)
         attn_pool_head: bool = self.config.get("attn_pool_head", False)
@@ -118,6 +120,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
         self.num_reg_tokens = num_reg_tokens
         self.attn_pool_special_tokens = attn_pool_special_tokens
         self.norm_layer = norm_layer
+        self.norm_layer_eps = norm_layer_eps
         self.mlp_layer = mlp_layer
         self.act_layer = act_layer
         self.rope_rot_type = rope_rot_type
@@ -190,6 +193,8 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             attention_dropout,
             dpr,
             pre_norm=pre_norm,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
             activation_layer=act_layer,
             layer_scale_init_value=layer_scale_init_value,
             norm_layer=norm_layer,
@@ -231,6 +236,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             rope_temperature=rope_temperature,
             layer_scale_init_value=layer_scale_init_value,
             norm_layer=norm_layer,
+            norm_layer_eps=norm_layer_eps,
             mlp_layer=mlp_layer,
             rope_rot_type=rope_rot_type,
         )
@@ -588,6 +594,7 @@ class RoPE_FlexiViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin,
             rope_temperature=self.rope_temperature,
             layer_scale_init_value=self.layer_scale_init_value,
             norm_layer=self.norm_layer,
+            norm_layer_eps=self.norm_layer_eps,
             mlp_layer=self.mlp_layer,
             rope_rot_type=self.rope_rot_type,
         )

birder/net/rope_vit.py

@@ -150,6 +150,10 @@ class RoPEAttention(nn.Module):
         attn_drop: float,
         proj_drop: float,
         num_special_tokens: int,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        norm_layer_eps: float = 1e-6,
         rope_rot_type: str = "standard",
     ) -> None:
         super().__init__()
@@ -167,7 +171,14 @@ class RoPEAttention(nn.Module):
         else:
             raise ValueError(f"Unknown rope_rot_type, got '{rope_rot_type}'")
 
-        self.qkv = nn.Linear(dim, dim * 3)
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        if qk_norm is True:
+            self.q_norm = norm_layer(self.head_dim, eps=norm_layer_eps)
+            self.k_norm = norm_layer(self.head_dim, eps=norm_layer_eps)
+        else:
+            self.q_norm = nn.Identity()
+            self.k_norm = nn.Identity()
+
         self.attn_drop = nn.Dropout(attn_drop)
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
@@ -176,6 +187,8 @@ class RoPEAttention(nn.Module):
         (B, N, C) = x.size()
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
         (q, k, v) = qkv.unbind(0)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
 
         n = self.num_special_tokens
         q = torch.concat([q[:, :, :n, :], self.apply_rot_fn(q[:, :, n:, :], rope)], dim=2)
@@ -207,6 +220,8 @@ class EncoderBlock(nn.Module):
         norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
         norm_layer_eps: float = 1e-6,
         mlp_layer: Callable[..., nn.Module] = FFN,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
         rope_rot_type: str = "standard",
     ) -> None:
         super().__init__()
@@ -222,6 +237,10 @@ class EncoderBlock(nn.Module):
             attn_drop=attention_dropout,
             proj_drop=dropout,
             num_special_tokens=num_special_tokens,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            norm_layer=norm_layer,
+            norm_layer_eps=norm_layer_eps,
             rope_rot_type=rope_rot_type,
         )
         if layer_scale_init_value is not None:
@@ -249,7 +268,6 @@ class EncoderBlock(nn.Module):
 
 
 class Encoder(nn.Module):
-    # pylint: disable=too-many-arguments,too-many-positional-arguments
     def __init__(
         self,
         num_layers: int,
@@ -261,6 +279,8 @@ class Encoder(nn.Module):
         attention_dropout: float,
         dpr: list[float],
         pre_norm: bool = False,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
         activation_layer: Callable[..., nn.Module] = nn.GELU,
         layer_scale_init_value: Optional[float] = None,
         norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
@@ -293,6 +313,8 @@ class Encoder(nn.Module):
                     norm_layer=norm_layer,
                     norm_layer_eps=norm_layer_eps,
                     mlp_layer=mlp_layer,
+                    qkv_bias=qkv_bias,
+                    qk_norm=qk_norm,
                     rope_rot_type=rope_rot_type,
                 )
             )
@@ -331,6 +353,7 @@ class MAEDecoderBlock(nn.Module):
         rope_temperature: float,
         layer_scale_init_value: Optional[float] = None,
         norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        norm_layer_eps: float = 1e-6,
         mlp_layer: Callable[..., nn.Module] = FFN,
         rope_rot_type: str = "standard",
     ) -> None:
@@ -346,7 +369,7 @@ class MAEDecoderBlock(nn.Module):
         )
 
         # Attention block
-        self.norm1 = norm_layer(hidden_dim, eps=1e-6)
+        self.norm1 = norm_layer(hidden_dim, eps=norm_layer_eps)
         self.attn = RoPEAttention(
             hidden_dim,
             num_heads,
@@ -361,7 +384,7 @@ class MAEDecoderBlock(nn.Module):
             self.layer_scale_1 = nn.Identity()
 
         # MLP block
-        self.norm2 = norm_layer(hidden_dim, eps=1e-6)
+        self.norm2 = norm_layer(hidden_dim, eps=norm_layer_eps)
         self.mlp = mlp_layer(hidden_dim, mlp_dim, act_layer=activation_layer, dropout=0.0)
         if layer_scale_init_value is not None:
             self.layer_scale_2 = LayerScale(hidden_dim, layer_scale_init_value)
@@ -403,6 +426,8 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         layer_scale_init_value: Optional[float] = self.config.get("layer_scale_init_value", None)
         pre_norm: bool = self.config.get("pre_norm", False)
         post_norm: bool = self.config.get("post_norm", True)
+        qkv_bias: bool = self.config.get("qkv_bias", True)
+        qk_norm: bool = self.config.get("qk_norm", False)
         num_reg_tokens: int = self.config.get("num_reg_tokens", 0)
         class_token: bool = self.config.get("class_token", True)
         attn_pool_head: bool = self.config.get("attn_pool_head", False)
@@ -450,6 +475,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
         self.num_reg_tokens = num_reg_tokens
         self.attn_pool_special_tokens = attn_pool_special_tokens
         self.norm_layer = norm_layer
+        self.norm_layer_eps = norm_layer_eps
         self.mlp_layer = mlp_layer
         self.act_layer = act_layer
         self.rope_rot_type = rope_rot_type
@@ -521,6 +547,8 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
             attention_dropout,
             dpr,
             pre_norm=pre_norm,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
             activation_layer=act_layer,
             layer_scale_init_value=layer_scale_init_value,
             norm_layer=norm_layer,
@@ -562,6 +590,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
             rope_temperature=rope_temperature,
             layer_scale_init_value=layer_scale_init_value,
             norm_layer=norm_layer,
+            norm_layer_eps=norm_layer_eps,
             mlp_layer=mlp_layer,
             rope_rot_type=rope_rot_type,
         )
@@ -904,6 +933,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
             rope_temperature=self.rope_temperature,
             layer_scale_init_value=self.layer_scale_init_value,
             norm_layer=self.norm_layer,
+            norm_layer_eps=self.norm_layer_eps,
             mlp_layer=self.mlp_layer,
             rope_rot_type=self.rope_rot_type,
         )
@@ -931,6 +961,7 @@ class RoPE_ViT(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, Mask
 #     - rms         : RMSNorm (instead of LayerNorm)
 #     - pn          : Pre-Norm (layer norm before the encoder) - implies different norm eps
 #     - npn         : No Post Norm (disables post-normalization layer)
+#     - qkn         : QK Norm
 #
 #     Feed-Forward Network:
 #     - swiglu      : SwiGLU FFN layer type (instead of standard FFN)
@@ -1068,6 +1099,20 @@ registry.register_model_config(
         "drop_path_rate": 0.1,
     },
 )
+registry.register_model_config(
+    "rope_vit_b16_qkn_ls",
+    RoPE_ViT,
+    config={
+        "patch_size": 16,
+        "num_layers": 12,
+        "num_heads": 12,
+        "hidden_dim": 768,
+        "mlp_dim": 3072,
+        "layer_scale_init_value": 1e-5,
+        "qk_norm": True,
+        "drop_path_rate": 0.1,
+    },
+)
 registry.register_model_config(
     "rope_i_vit_b16_pn_aps_c1",  # For PE Core - https://arxiv.org/abs/2504.13181
     RoPE_ViT,

birder/net/smt.py

@@ -275,7 +275,6 @@ class Stem(nn.Module):
 
 
 class SMTStage(nn.Module):
-    # pylint: disable=too-many-arguments,too-many-positional-arguments
     def __init__(
         self,
         dim: int,

birder/net/ssl/ibot.py

@@ -25,7 +25,6 @@ from birder.net.ssl.dino_v1 import DINOHead
 
 # pylint: disable=invalid-name
 class iBOTLoss(nn.Module):
-    # pylint: disable=too-many-arguments,too-many-positional-arguments
     def __init__(
         self,
         out_dim: int,