(self, x)
| 161 | self.linear = nn.Linear(in_channels * downscaling_factor ** 2, out_channels) |
| 162 | |
| 163 | def forward(self, x): |
| 164 | b, c, h, w = x.shape |
| 165 | new_h, new_w = h // self.downscaling_factor, w // self.downscaling_factor |
| 166 | x = self.patch_merge(x).view(b, -1, new_h, new_w).permute(0, 2, 3, 1) |
| 167 | x = self.linear(x) |
| 168 | return x |
| 169 | |
| 170 | |
| 171 | class StageModule(nn.Module): |