def sync_workers():
    """
    Yields distributed rank and synchronizes all workers on exit.
    """
    rank = get_rank()
    yield rank
    barrier()

    Yields distributed rank and synchronizes all workers on exit.
    
training/src/utils/distributed.py