Bull

pip install einops from einops import rearrange x = torch.randn(64, 3, 32, 32) patch_size = 4 # 4 pixels print('x :', x.shape) patches = rearrange(x, 'b c (h s1) (w s2) -> b (h w) (s1 s2 c)', s1=patch_size, s2=patch_size) # 64 3 32 32 = 64 3 8*4 8*4 -> 64 8*8 4*4*3 print('patches :', patches.shape) x : torch.Size([64, 3, 32, 32]) patches : torch.Size([64, 64, 48]) rearrange(x, 'b c (h s1) (w s2)..