It looks like your code snippet got cut off at the end of the `PatchEmbed` class definition. Let me help you complete it and provide a clean, fully functional version of the `PatchEmbed` module from the original Vision Transformer (ViT) implementation by rwightman, as referenced in your comment. Here’s the **complete and corrected** `PatchEmbed` class: ```python class PatchEmbed(nn.Module): """ 2D Image to Patch Embedding """ def __init__(self, img_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=None): super().__init__() img_size = (img_size, img_size) patch_size = (patch_size, patch_size) self.img_size = img_size self.patch_size = patch_size self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) self.num_patches = self.grid_size[0] * self.grid_size[1] # Convolutional layer to project image patches into embedding space self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size) # Optional normalization layer after projection self.norm = norm_layer(embed_dim) if norm_layer is not None else nn.Identity() def forward(self, x): B, C, H, W = x.shape assert H == self.img_size[0] and W == self.img_size[1], \ f"Input image size ({H}x{W}) doesn't match model ({self.img_size[0]}x{self.img_size[1]})." # [B, C, H, W] -> [B, embed_dim, num_patches_h, num_patches_w] x = self.proj(x) # Shape: (B, embed_dim, H//P, W//P) # Flatten spatial dimensions: [B, embed_dim, num_patches_h, num_patches_w] -> [B, embed_dim, num_patches] x = x.flatten(2).transpose(1, 2) # Shape: (B, num_patches, embed_dim) # Apply normalization if specified x = self.norm(x) return x ``` --- ### ✅ Summary of Key Features: - **`img_size`**: Input image size (e.g., 224×224). - **`patch_size`**: Size of each patch (e.g., 16×16 → 196 patches for 224×224). - **`in_c`**: Number of input channels (usually 3 for RGB). - **`embed_dim`**: Dimensionality of the output embedding (e.g., 768). - **`norm_layer`**: Optional normalization (like LayerNorm), defaults to identity. - **Output shape**: `(B, num_patches, embed_dim)` — ready for transformer input. --- ### 🔧 Example Usage: ```python # Example: 224x224 image, 16x16 patches, 3 input channels, 768-dim embeddings patch_embed = PatchEmbed(img_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=nn.LayerNorm) x = torch.randn(1, 3, 224, 224) # Batch size 1 out = patch_embed(x) print(out.shape) # Output: torch.Size([1, 196, 768]) ``` > 💡 There are 196 patches: (224 / 16) × (224 / 16) = 14 × 14 = 196. --- Let me know if you'd like the full Vision Transformer architecture built on top of this (`VisionTransformer` class), or want to add positional embeddings, attention blocks, etc.