25 lines
716 B
Python
25 lines
716 B
Python
"""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)
|