diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py
index e98f455..9a34d56 100644
--- a/dags/test-train-pytorch.py
+++ b/dags/test-train-pytorch.py
@@ -81,11 +81,13 @@ def run_training_node_func(rank, world_size):
   my_fqdn = socket.getfqdn()
   print(f"[{rank}] My FQDN: {my_fqdn}")
 
+  master_addr = None  # Initialize outside if/else
   if rank == 0:
       # Rank 0 stores its address as the master
       sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}'))
       sync_state['master_addr'] = my_fqdn
       Variable.set('ddp_sync_state', json.dumps(sync_state))
+      master_addr = my_fqdn  # For rank 0, master is itself
       print(f"[{rank}] I am the master. Stored my address: {my_fqdn}")
   else:
       # Other ranks wait for rank 0 to store its address
@@ -106,9 +108,9 @@ def run_training_node_func(rank, world_size):
   time.sleep(3)
 
   # STEP 3: Configure distributed environment
-  # Use dynamic master address from rank 0
-  actual_master_addr = my_fqdn if rank == 0 else sync_state.get('master_addr', MASTER_ADDR)
-  os.environ['MASTER_ADDR'] = actual_master_addr
+  # Use the master_addr that was determined above
+  print(f"[{rank}] Using master address: {master_addr}")
+  os.environ['MASTER_ADDR'] = master_addr
   os.environ['MASTER_PORT'] = MASTER_PORT
   os.environ['WORLD_SIZE'] = str(world_size)
   os.environ['RANK'] = str(rank)