    def forward_features(self, data: torch.Tensor) -> torch.Tensor:
        """
        Forward pass of feature extraction.
        Args:
            data (`torch.Tensor`):
                Input images of the shape [B, C, H, W].
        Returns:
            output (`torch.Tensor`):
                Image features of the backbone.
        """
        output = data
        for stage in self.stages:
            output = stage(output)
        return output
Forward pass of feature extraction.
Args:
data (`torch.Tensor`):
Input images of the shape [B, C, H, W].
Returns:
output (`torch.Tensor`):
Image features of the backbone.
towhee/models/max_vit/max_vit.py