(self, dim)
| 239 | """ |
| 240 | |
| 241 | def __init__(self, dim): |
| 242 | super().__init__() |
| 243 | self.dim = dim |
| 244 | |
| 245 | # layers |
| 246 | self.norm = RMS_norm(dim) |
| 247 | self.to_qkv = nn.Conv2d(dim, dim * 3, 1) |
| 248 | self.proj = nn.Conv2d(dim, dim, 1) |
| 249 | |
| 250 | # zero out the last layer params |
| 251 | nn.init.zeros_(self.proj.weight) |
| 252 | |
| 253 | def forward(self, x): |
| 254 | identity = x |