diffusion_models.models.nn.checkpoint¶
- diffusion_models.models.nn.checkpoint(func, inputs, params, flag)[source]¶
Evaluate a function without caching intermediate activations, allowing for reduced memory at the expense of extra compute in the backward pass.
- Parameters:
func – the function to evaluate.
inputs – the argument sequence to pass to func.
params – a sequence of parameters func depends on but does not explicitly take as arguments.
flag – if False, disable gradient checkpointing.