checkpoint_sequential

dragon.vm.torch.utils.checkpoint.checkpoint_sequential(
  functions,
  input,
  segments=1,
  **kwargs
)[source]

Apply functions and create segmental checkpoints.

Parameters:
  • functions (Union[torch.nn.Sequential, Sequence[callable]]) – The functions to apply sequentially.
  • input (dragon.vm.torch.Tensor) – The input tensor.
  • segments (Union[int, Sequence[int]], optional) – The number or size of chunked checkpoints.
Returns:

Any – The function outputs.