def forward(self, x: Tensor) -> Tensor:
        """Classify images if ``mobile_sam=False``, produce feature maps if ``mobile_sam=True``."""
        x = self.patch_embed(x)
        x = self.layers(x)

        if self.mobile_sam:
            # MobileSAM
            x = x.unflatten(1, (self.feat_size, self.feat_size)).permute(0, 3, 1, 2)
            x = self.neck(x)  # type: ignore
        else:
            # classification
            x = x.mean(1)
            x = self.head(self.norm_head(x))
        return x
Classify images if ``mobile_sam=False``, produce feature maps if ``mobile_sam=True``.
kornia/contrib/models/tiny_vit.py