orbit.utils.LoopFnWithSummaries

Implements a two-program approach for optimizing summaries on TPU.

Inherits From: OptionalSummariesFunction

orbit.utils.LoopFnWithSummaries(
    function, **tf_function_kwargs
)

This version works with the result of create_tf_while_loop_fn.

function The underlying function to wrap.
**tf_function_kwargs Additional arguments to pass to tf.function.

Methods

__call__

View source

__call__(
    iterator, num_steps
)

Call self as a function.