View source on GitHub |
Creates a loop function compatible with TF's AutoGraph loop conversion.
orbit.utils.create_tf_while_loop_fn(
step_fn
)
Args | |
---|---|
step_fn
|
A function taking a nested structure of tf.data.Iterator or
DistributedIterator . Currently, any return values are ignored.
|
Returns | |
---|---|
A loop function taking required iterator and num_steps parameters. If
called inside a tf.function , the loop will be converted by AutoGraph into
a tf.while_loop construct. See the loop_fn definition below for
additional details.
|