(self, x, condition=None)
| 394 | self.grad_checkpointing = flag |
| 395 | |
| 396 | def forward(self, x, condition=None): |
| 397 | for res_i, module in enumerate(self.resblocks): |
| 398 | if self.grad_checkpointing: |
| 399 | x = checkpoint(module, x, res_i, condition, use_reentrant=False) |
| 400 | else: |
| 401 | x = module(x, res_i, condition) |
| 402 | |
| 403 | return x |
| 404 | |
| 405 | class TransformerBase(nn.Module): |
| 406 | def __init__(self, width, layers, heads, window_size, token_len, block_cls, drop_path_rate=0.0): |
nothing calls this directly
no outgoing calls
no test coverage detected