Wrapper for MLPerf compliance logging calls.
    All arguments but 'sync' are passed to mlperf_log.gnmt_print function.
    If 'sync' is set to True then the wrapper will synchronize all distributed
    workers. 'sync' should be set to True for all compliance tags that require
    accurate timing (RUN_START, RUN_STOP etc.)

def gnmt_print(*args, **kwargs):
    """
    Wrapper for MLPerf compliance logging calls.
    All arguments but 'sync' are passed to mlperf_log.gnmt_print function.
    If 'sync' is set to True then the wrapper will synchronize all distributed
    workers. 'sync' should be set to True for all compliance tags that require
    accurate timing (RUN_START, RUN_STOP etc.)
    """
    if kwargs.pop('sync'):
        barrier()
    if get_rank() == 0:
        kwargs['stack_offset'] = 2
        mlperf_log.gnmt_print(*args, **kwargs)