(self, hidden_states)
| 44 | |
| 45 | |
| 46 | def patchify(self, hidden_states): |
| 47 | hidden_states = rearrange(hidden_states, "B C (H P) (W Q) -> B (H W) (C P Q)", P=2, Q=2) |
| 48 | return hidden_states |
| 49 | |
| 50 | |
| 51 | def align_res_stack_to_original_blocks(self, res_stack, num_blocks, hidden_states): |