(self, x, return_attention=False)
| 233 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 234 | |
| 235 | def forward(self, x, return_attention=False): |
| 236 | y, attn = self.attn(self.norm1(x)) |
| 237 | if return_attention: |
| 238 | return attn |
| 239 | x = x + self.drop_path(y) |
| 240 | x = x + self.drop_path(self.mlp(self.norm2(x))) |
| 241 | return x |
| 242 | |
| 243 | |
| 244 | class PatchEmbed(nn.Module): |