Function get_local_rank

PATH/core/distributed_utils.py:163–164

()

    return int(os.environ.get('SLURM_PROCID', 0))

def get_local_rank():
    return get_rank() % torch.cuda.device_count()

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