| 199 | |
| 200 | |
| 201 | class Transformer(nn.Module): |
| 202 | def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None): |
| 203 | super().__init__() |
| 204 | self.width = width |
| 205 | self.layers = layers |
| 206 | self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)]) |
| 207 | |
| 208 | def forward(self, x: torch.Tensor): |
| 209 | return self.resblocks(x) |
| 210 | |
| 211 | |
| 212 | class VisualTransformer(nn.Module): |