(self, x: torch.Tensor, grid_size: torch.Tensor)
| 584 | return x, grid_size # x, grid_size: (f, h, w) |
| 585 | |
| 586 | def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor): |
| 587 | return rearrange( |
| 588 | x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)', |
| 589 | f=grid_size[0], h=grid_size[1], w=grid_size[2], |
| 590 | x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2] |
| 591 | ) |
| 592 | |
| 593 | def forward(self, |
| 594 | x: torch.Tensor, |