(self, x, relative_pos_enc=None)
| 489 | drop_path) if drop_path > 0. else nn.Identity() |
| 490 | |
| 491 | def _forward_impl(self, x, relative_pos_enc=None): |
| 492 | # print(x.shape) |
| 493 | # x = x + self.pos_embed(x) |
| 494 | x = x + self.drop_path(self.token_mixer(self.norm1(x), relative_pos_enc)) |
| 495 | x = x + self.drop_path(self.mlp(self.norm2(x))) |
| 496 | return x |
| 497 | |
| 498 | def forward(self, x, relative_pos_enc=None): |
| 499 | if self.grad_checkpoint and x.requires_grad: |