patch.py

from itertools import repeat
import pdb
class PatchEmbed(nn.Module):
  def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
    super().__init__()
    img_size = tuple(repeat(img_size, 2))
    patch_size = tuple(repeat(patch_size, 2))
    num_patches = (img_size[1] // patch_size[1]) * (img_size[0] //   patch_size[0])
    self.img_size = img_size
    self.patch_size = patch_size
    self.num_patches = num_patches
    self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size,    stride=patch_size)
def forward(self, x):
    B, C, H, W = x.shape
    x = self.proj(x).flatten(2).transpose(1, 2)
    return x