birder 0.3.0__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/fs_ops.py

@@ -627,7 +627,7 @@ def load_model(
         net.to(dtype)
     if inference is True:
         for param in net.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         if pt2 is False:  # NOTE: Remove when GraphModule add support for 'eval'
             net.eval()
@@ -799,7 +799,7 @@ def load_detection_model(
         net.to(dtype)
     if inference is True:
         for param in net.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         net.eval()
 

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/common/training_cli.py

@@ -39,6 +39,7 @@ def add_optimization_args(parser: argparse.ArgumentParser, default_batch_size: i
     group = parser.add_argument_group("Optimization parameters")
     group.add_argument("--batch-size", type=int, default=default_batch_size, metavar="N", help="the batch size")
     group.add_argument("--opt", type=str, choices=list(get_args(OptimizerType)), default="sgd", help="optimizer to use")
+    group.add_argument("--opt-fused", default=False, action="store_true", help="use fused optimizer implementation")
     group.add_argument("--momentum", type=float, default=0.9, metavar="M", help="optimizer momentum")
     group.add_argument("--nesterov", default=False, action="store_true", help="use nesterov momentum")
     group.add_argument("--opt-eps", type=float, help="optimizer epsilon (None to use the optimizer default)")
@@ -249,6 +250,7 @@ def add_data_aug_args(
     default_level: int = 4,
     default_min_scale: Optional[float] = None,
     default_re_prob: Optional[float] = None,
+    smoothing_alpha: bool = False,
     mixup_cutmix: bool = False,
 ) -> None:
     group = parser.add_argument_group("Data augmentation parameters")
@@ -285,6 +287,8 @@ def add_data_aug_args(
     group.add_argument(
         "--simple-crop", default=False, action="store_true", help="use simple random crop (SRC) instead of RRC"
     )
+    if smoothing_alpha is True:
+        group.add_argument("--smoothing-alpha", type=float, default=0.0, help="label smoothing alpha")
     if mixup_cutmix is True:
         group.add_argument("--mixup-alpha", type=float, help="mixup alpha")
         group.add_argument("--cutmix", default=False, action="store_true", help="enable cutmix")
@@ -565,9 +569,9 @@ def add_training_data_args(parser: argparse.ArgumentParser, unsupervised: bool =
     group.add_argument("--wds", default=False, action="store_true", help="use webdataset for training")
     group.add_argument("--wds-info", type=str, metavar="FILE", help="wds info file path")
     group.add_argument("--wds-cache-dir", type=str, metavar="DIR", help="webdataset cache directory")
-    group.add_argument("--wds-train-size", type=int, metavar="N", help="size of the wds training set")
     if unsupervised is False:
         group.add_argument("--wds-class-file", type=str, metavar="FILE", help="class list file")
+        group.add_argument("--wds-train-size", type=int, metavar="N", help="size of the wds training set")
         group.add_argument("--wds-val-size", type=int, metavar="N", help="size of the wds validation set")
         group.add_argument(
             "--wds-training-split", type=str, default="training", metavar="NAME", help="wds dataset train split"
@@ -576,6 +580,7 @@ def add_training_data_args(parser: argparse.ArgumentParser, unsupervised: bool =
             "--wds-val-split", type=str, default="validation", metavar="NAME", help="wds dataset validation split"
         )
     else:
+        group.add_argument("--wds-size", type=int, metavar="N", help="size of the wds")
         group.add_argument(
             "--wds-split", type=str, default="training", metavar="NAME", help="wds dataset split to load"
         )

birder/common/training_utils.py

@@ -593,12 +593,14 @@ def get_optimizer(parameters: list[dict[str, Any]], l_rate: float, args: argpars
         kwargs["betas"] = args.opt_betas
     if getattr(args, "opt_alpha", None) is not None:
         kwargs["alpha"] = args.opt_alpha
+    if getattr(args, "opt_fused", False) is True:
+        kwargs["fused"] = True
 
     # For optimizer compilation
     # lr = torch.tensor(l_rate) - Causes weird LR scheduling bugs
     lr = l_rate
-    if getattr(args, "compile_opt", False) is not False:
-        if opt not in ("lamb", "lambw", "lars"):
+    if getattr(args, "compile_opt", False) is True:
+        if opt not in ("sgd", "lamb", "lambw", "lars"):
             logger.debug("Setting optimizer capturable to True")
             kwargs["capturable"] = True
 

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/base.py

@@ -173,14 +173,14 @@ class BaseNet(nn.Module):
 
     def freeze(self, freeze_classifier: bool = True, unfreeze_features: bool = False) -> None:
         for param in self.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         if freeze_classifier is False:
             for param in self.classifier.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
         if unfreeze_features is True and hasattr(self, "features") is True:
             for param in self.features.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         """

birder/net/biformer.py

@@ -468,14 +468,14 @@ class BiFormer(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/cas_vit.py

@@ -269,18 +269,18 @@ class CAS_ViT(DetectorBackbone):
 
     def freeze(self, freeze_classifier: bool = True, unfreeze_features: bool = False) -> None:
         for param in self.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         if freeze_classifier is False:
             for param in self.classifier.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
             for param in self.dist_classifier.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
         if unfreeze_features is True:
             for param in self.features.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
     def transform_to_backbone(self) -> None:
         self.features = nn.Identity()
@@ -300,14 +300,14 @@ class CAS_ViT(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/coat.py

@@ -563,24 +563,24 @@ class CoaT(DetectorBackbone):
     def freeze_stages(self, up_to_stage: int) -> None:
         if up_to_stage >= 1:
             for param in self.patch_embed1.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
             for param in self.serial_blocks1.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
         if up_to_stage >= 2:
             for param in self.patch_embed2.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
             for param in self.serial_blocks2.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
         if up_to_stage >= 3:
             for param in self.patch_embed3.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
             for param in self.serial_blocks3.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
         if up_to_stage >= 4:
             for param in self.patch_embed4.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
             for param in self.serial_blocks4.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         features = self._features(x)

birder/net/conv2former.py

@@ -218,14 +218,14 @@ class Conv2Former(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def masked_encoding_retention(
         self,

birder/net/convnext_v1.py

@@ -158,14 +158,14 @@ class ConvNeXt_v1(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def masked_encoding_retention(
         self,
@@ -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/convnext_v2.py

@@ -180,14 +180,14 @@ class ConvNeXt_v2(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def masked_encoding_retention(
         self,

birder/net/crossformer.py

@@ -404,14 +404,14 @@ class CrossFormer(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.patch_embed.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.patch_embed(x)

birder/net/cspnet.py

@@ -342,14 +342,14 @@ class CSPNet(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/cswin_transformer.py

@@ -359,14 +359,14 @@ class CSWin_Transformer(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/darknet.py

@@ -115,14 +115,14 @@ class Darknet(DetectorBackbone):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.stem(x)

birder/net/davit.py

@@ -391,14 +391,14 @@ class DaViT(DetectorBackbone, PreTrainEncoder, MaskedTokenRetentionMixin):
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.stem.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         for idx, module in enumerate(self.body.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def masked_encoding_retention(
         self,

birder/net/deit.py

@@ -117,14 +117,14 @@ class DeiT(BaseNet):
 
     def freeze(self, freeze_classifier: bool = True, unfreeze_features: bool = False) -> None:
         for param in self.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
         if freeze_classifier is False:
             for param in self.classifier.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
             for param in self.dist_classifier.parameters():
-                param.requires_grad = True
+                param.requires_grad_(True)
 
     def set_causal_attention(self, is_causal: bool = True) -> None:
         self.encoder.set_causal_attention(is_causal)

birder/net/deit3.py

@@ -182,16 +182,16 @@ class DeiT3(DetectorBackbone, PreTrainEncoder, MaskedTokenOmissionMixin, MaskedT
 
     def freeze_stages(self, up_to_stage: int) -> None:
         for param in self.conv_proj.parameters():
-            param.requires_grad = False
+            param.requires_grad_(False)
 
-        self.pos_embedding.requires_grad = False
+        self.pos_embedding.requires_grad_(False)
 
         for idx, module in enumerate(self.encoder.children()):
             if idx >= up_to_stage:
                 break
 
             for param in module.parameters():
-                param.requires_grad = False
+                param.requires_grad_(False)
 
     def set_causal_attention(self, is_causal: bool = True) -> None:
         self.encoder.set_causal_attention(is_causal)