(self, width, layers, heads, window_size, token_len, block_cls, drop_path_rate=0.0)
| 404 | |
| 405 | class TransformerBase(nn.Module): |
| 406 | def __init__(self, width, layers, heads, window_size, token_len, block_cls, drop_path_rate=0.0): |
| 407 | super().__init__() |
| 408 | self.layernorm1 = LayerNorm(width) |
| 409 | self.transformer = Transformer(width, layers, heads, window_size=window_size, block_cls=block_cls, drop_path_rate=drop_path_rate) |
| 410 | self.layernorm2 = LayerNorm(width) |
| 411 | |
| 412 | def set_grad_checkpointing(self, set_checkpointing=True): |
| 413 | self.transformer.set_grad_checkpointing(set_checkpointing) |
nothing calls this directly
no test coverage detected