(x, patch_size)
| 302 | |
| 303 | |
| 304 | def unpatchify(x, patch_size): |
| 305 | if patch_size == 1: |
| 306 | return x |
| 307 | |
| 308 | if x.dim() == 4: |
| 309 | x = rearrange( |
| 310 | x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) |
| 311 | elif x.dim() == 5: |
| 312 | x = rearrange( |
| 313 | x, |
| 314 | "b (c r q) f h w -> b c f (h q) (w r)", |
| 315 | q=patch_size, |
| 316 | r=patch_size, |
| 317 | ) |
| 318 | return x |
| 319 | |
| 320 | |
| 321 | class AvgDown3D(nn.Module): |