(self, x)
| 29 | self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() |
| 30 | |
| 31 | def forward(self, x): |
| 32 | B, C, H, W = x.shape |
| 33 | _assert(H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).") |
| 34 | _assert(W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).") |
| 35 | x = self.proj(x) |
| 36 | if self.flatten: |
| 37 | x = x.flatten(2).transpose(1, 2) # BCHW -> BNC |
| 38 | x = self.norm(x) |
| 39 | return x |