Function get_rank

PATH/core/distributed_utils.py:160–161

()

                raise RuntimeError('buffer {} does not have task_specific'.format(name, id(b)))

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

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