| """Utility functions for training. |
| |
| For licensing see accompanying LICENSE file. |
| Copyright (C) 2025 Apple Inc. All Rights Reserved. |
| """ |
|
|
| from torch.utils.checkpoint import checkpoint |
|
|
|
|
| def checkpoint_wrapper(self, fn, *args): |
| """Helper function that applies checkpointing. |
| |
| If enabled applies grad checkpointing, otherwise just executes the function normally. |
| """ |
| if not hasattr(self, "grad_checkpointing"): |
| raise AttributeError( |
| "Trying to apply grad checkpointing on a model that does not have a grad_checkpointing " |
| "attribute." |
| ) |
|
|
| if self.grad_checkpointing: |
| return checkpoint(fn, *args, use_reentrant=False) |
| else: |
| return fn(*args) |
|
|
