(self, x)
| 119 | self.register_buffer("attn_mask", attn_mask) |
| 120 | |
| 121 | def forward(self, x): |
| 122 | |
| 123 | shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(2, 3)) |
| 124 | sw_x, _, _ = super().forward(shifted_x, self.attn_mask) |
| 125 | x = torch.roll(sw_x, shifts=(self.shift_size, self.shift_size), dims=(2, 3)) |
| 126 | |
| 127 | return x, None, None |
| 128 | |
| 129 | |
| 130 | class DAttentionBaseline(nn.Module): |