(self, x, H, W)
| 286 | self.downsample = None |
| 287 | |
| 288 | def forward(self, x, H, W): |
| 289 | for blk in self.blocks: |
| 290 | blk.H, blk.W = H, W |
| 291 | if self.use_checkpoint: |
| 292 | x = checkpoint.checkpoint(blk, x) |
| 293 | else: |
| 294 | x = blk(x) |
| 295 | |
| 296 | if self.downsample is not None: |
| 297 | x = x.transpose(1, 2).reshape(x.shape[0], -1, H, W) |
| 298 | x, Ho, Wo = self.downsample(x) |
| 299 | else: |
| 300 | Ho, Wo = H, W |
| 301 | return x, Ho, Wo |
| 302 | |
| 303 | def extra_repr(self) -> str: |
| 304 | return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}" |
nothing calls this directly
no outgoing calls
no test coverage detected