Helper function that applies checkpointing. If enabled applies grad checkpointing, otherwise just executes the function normally.
(self, fn, *args)
| 8 | |
| 9 | |
| 10 | def checkpoint_wrapper(self, fn, *args): |
| 11 | """Helper function that applies checkpointing. |
| 12 | |
| 13 | If enabled applies grad checkpointing, otherwise just executes the function normally. |
| 14 | """ |
| 15 | if not hasattr(self, "grad_checkpointing"): |
| 16 | raise AttributeError( |
| 17 | "Trying to apply grad checkpointing on a model that does not have a grad_checkpointing " |
| 18 | "attribute." |
| 19 | ) |
| 20 | |
| 21 | if self.grad_checkpointing: |
| 22 | return checkpoint(fn, *args, use_reentrant=False) |
| 23 | else: |
| 24 | return fn(*args) |