torch.autograd.save_on_cpu
function save_on_cpu<T>(fn: () => T, _options?: { pin_memory?: boolean }): TContext manager to save tensors on CPU during forward pass.
This is a convenience wrapper around saved_tensors_hooks that automatically moves tensors to CPU when saved and back to the original device when unpacked.
This can significantly reduce GPU memory usage during training at the cost of CPU-GPU transfer overhead during backward.
Parameters
fn() => T- Function to execute with save-on-CPU enabled
_options{ pin_memory?: boolean }optional
Returns
T– The result of the functionExamples
// Save GPU memory by storing activations on CPU
torch.autograd.graph.save_on_cpu(() => {
const y = model.forward(x);
y.backward();
});