recompute

paddle.distributed.fleet.utils. recompute ( function, *args, **kwargs ) [source]

recompute intermediate activations to save then memory.

Parameters
  • function – layer of sequence of layers that describes part of forward pass of the model whose

  • activations will be released to save memory in forward stage and will be recomputed (intermediate) –

  • backward stage for gradient calculation. (in) –

  • preserve_rng_state (bool, optional) – if preserve the RNG state of forward and restore it in backward.

  • args – inputs to the function

Returns

Output of function on args