Function get_world_size

PATH/core/distributed_utils.py:166–167

()

    return get_rank() % torch.cuda.device_count()

def get_world_size():
    return int(os.environ.get('SLURM_NTASKS', 1))

def dist_init(method='slurm', port='5671'):
    assert method == 'slurm'