def patchify(self, imgs):
        """Split images into non-overlapped patches.

        Args:
            imgs (torch.Tensor): A batch of images, of shape B x C x H x W.
        Returns:
            torch.Tensor: Patchified images. The shape is B x L x D.
        """
        p = self.patch_size
        assert len(imgs.shape
                   ) == 4 and imgs.shape[2] % p == 0 and imgs.shape[3] % p == 0

        B, C, ori_h, ori_w = imgs.shape
        h = ori_h // p
        w = ori_w // p
        x = imgs.reshape(shape=(B, C, h, p, w, p))
        x = torch.einsum('bchpwq->bhwpqc', x)

        # (B, f*f, downsample_raito*downsample_raito*3)
        x = x.reshape(shape=(B, h * w, p**2 * C))
        return x
Split images into non-overlapped patches.

        Args:
            imgs (torch.Tensor): A batch of images, of shape B x C x H x W.
        Returns:
            torch.Tensor: Patchified images. The shape is B x L x D.
        
mmpretrain/models/heads/spark_head.py