Method wait_for_block

utils/offloading.py:277–283

(self, block_idx: int)

        clean_memory_on_device(self.device)

    def wait_for_block(self, block_idx: int):
        if self.blocks_to_swap is None or self.blocks_to_swap == 0:
            return
        if self.reentrant_activation_checkpointing and torch.is_grad_enabled():
            # Second forward pass, don't do block swapping
            return
        self._wait_blocks_move(block_idx)

    def submit_move_blocks_forward(self, block_idx: int):
        # check if blocks_to_swap is enabled