    def forward(self, data: torch.Tensor) -> torch.Tensor:
        """ Forward pass
        Args:
            data (`torch.Tensor`):
                Input images of the shape [B, C, H, W].
        Returns:
            output (`torch.Tensor`):
                Classification output of the shape [B, num_classes].
        """
        output = self.forward_features(self.stem(data))
        output = self.forward_head(output)
        return output
Forward pass
Args:
data (`torch.Tensor`):
Input images of the shape [B, C, H, W].
Returns:
output (`torch.Tensor`):
Classification output of the shape [B, num_classes].
towhee/models/max_vit/max_vit.py