| 217 | """ |
| 218 | # |
| 219 | def __init__(self, in_channels): |
| 220 | super().__init__() |
| 221 | self.in_channels = in_channels |
| 222 | |
| 223 | self.norm = Normalize(in_channels) |
| 224 | self.q = torch.nn.Conv2d(in_channels, |
| 225 | in_channels, |
| 226 | kernel_size=1, |
| 227 | stride=1, |
| 228 | padding=0) |
| 229 | self.k = torch.nn.Conv2d(in_channels, |
| 230 | in_channels, |
| 231 | kernel_size=1, |
| 232 | stride=1, |
| 233 | padding=0) |
| 234 | self.v = torch.nn.Conv2d(in_channels, |
| 235 | in_channels, |
| 236 | kernel_size=1, |
| 237 | stride=1, |
| 238 | padding=0) |
| 239 | self.proj_out = torch.nn.Conv2d(in_channels, |
| 240 | in_channels, |
| 241 | kernel_size=1, |
| 242 | stride=1, |
| 243 | padding=0) |
| 244 | self.attention_op: Optional[Any] = None |
| 245 | |
| 246 | def forward(self, x): |
| 247 | h_ = x |