diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py
index 720a87f..497d4fd 100644
--- a/dags/test-train-pytorch.py
+++ b/dags/test-train-pytorch.py
@@ -84,6 +84,7 @@ def run_training_node_func(rank, world_size):
   os.environ['WORLD_SIZE'] = str(world_size)
   os.environ['RANK'] = str(rank)
   os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'
+  os.environ['NCCL_PORT_RANGE'] = '30000-30100'
   os.environ['NCCL_DEBUG'] = 'INFO'
   os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
   os.environ['NCCL_BLOCKING_WAIT'] = '1'