DeMemWM / utils /distributed_utils.py
BonanDing's picture
Add DeMemWM curriculum training support
b165569
Raw
History Blame Contribute Delete
150 Bytes
import os
def process_rank() -> int:
return int(os.environ.get("RANK", os.environ.get("SLURM_PROCID", 0)))
is_rank_zero = process_rank() == 0