Function cleanup

train.py:39–40

()

    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def save_checkpoint(args, model, optimizer=None, step=None, epoch=None):
    CHECKPOINT_PATH = None