(module, grad_input, grad_output)
| 243 | block_idx_to_wait = block_index - 1 |
| 244 | |
| 245 | def backward_hook(module, grad_input, grad_output): |
| 246 | if self.debug: |
| 247 | print(f"Backward hook for block {block_index}") |
| 248 | |
| 249 | if swapping: |
| 250 | self._submit_move_blocks(block_idx_to_cpu, block_idx_to_cuda) |
| 251 | if waiting: |
| 252 | self._wait_blocks_move(block_idx_to_wait) |
| 253 | return None |
| 254 | |
| 255 | return backward_hook |
| 256 |
nothing calls this directly
no test coverage detected