File size: 293 Bytes
c29babb
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
import os

# See more: https://pytorch.org/docs/stable/elastic/run.html#environment-variables
RANK = int(os.getenv("LOCAL_RANK", 0))
IS_GLOBAL_ZERO = RANK == 0
NODE_RANK = int(os.getenv("NODE_RANK", 0))
LOCAL_RANK = int(os.getenv("LOCAL_RANK", 0))
WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))