    def forward_head(self, data: torch.Tensor, pre_logits: bool = False):
        """
        Forward pass of classification head.
        Args:
            data (`torch.Tensor`):
                Input features.
            pre_logits (`bool`):
                If true pre-logits are returned.
        Returns:
            output (`torch.Tensor`):
                Classification output of the shape [B, num_classes].
        """
        if self.global_pool == "avg":
            data = data.mean(dim=(2, 3))
        elif self.global_pool == "max":
            data = torch.amax(data, dim=(2, 3))
        return data if pre_logits else self.head(data)
Forward pass of classification head.
Args:
data (`torch.Tensor`):
Input features.
pre_logits (`bool`):
If true pre-logits are returned.
Returns:
output (`torch.Tensor`):
Classification output of the shape [B, num_classes].
towhee/models/max_vit/max_vit.py