def mark_step():
torch_xla._XLAC._xla_step_marker(torch_xla._XLAC._xla_get_default_device(),
[],
wait=xu.getenv_as('XLA_SYNC_WAIT', bool,
False))
# Only emit metrics from the first local device index, to avoid emitting the
# same values from different threads.
if is_master_ordinal():
ms.save_metrics()
def _mark_step(replication):
devices = []
if replication:
replication.enter()
devices = replication.replication_devices()
torch_xla._XLAC._xla_step_marker(
torch_xla._XLAC._xla_get_default_device(), devices, wait=False)
# Only emit metrics from the first local device index, to avoid emitting the
# same values from different threads.
if getattr(_TLS, 'device_index', 0) == 0:
ms.save_metrics()