def forward(
        self, x: torch.Tensor
    ) -> tuple[list[torch.Tensor], list[Optional[torch.Tensor]]]:
        """
        Forward pass to extract multi-stage features.

        Args:
            x (torch.Tensor): Input tensor of shape (B, C, H, W).

        Returns:
            tuple[list[torch.Tensor], list[torch.Tensor]]: Tuple of feature maps and cls tokens (if supported) at different scales.
        """

        if self.has_prefix_tokens:
            features, prefix_tokens = self._forward_with_prefix_tokens(x)
        else:
            features = self._forward_without_prefix_tokens(x)
            prefix_tokens = [None] * len(features)

        return features, prefix_tokens

        Forward pass to extract multi-stage features.

        Args:
            x (torch.Tensor): Input tensor of shape (B, C, H, W).

        Returns:
            tuple[list[torch.Tensor], list[torch.Tensor]]: Tuple of feature maps and cls tokens (if supported) at different scales.
        
segmentation_models_pytorch/encoders/timm_vit.py