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.