diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py
index 612445d..853bd55 100644
--- a/dags/test-train-pytorch.py
+++ b/dags/test-train-pytorch.py
@@ -43,15 +43,16 @@ def run_training_node_func(rank, world_size):
   # This ensures NCCL reads these settings during initialization
   os.environ['NCCL_DEBUG'] = 'INFO'
   os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
-  os.environ['NCCL_BLOCKING_WAIT'] = '1'
-  os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'  # Force ethernet (reliable)
-  os.environ['NCCL_IB_DISABLE'] = '1'  # Disable IB - infrastructure not configured
+  os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1'  # Updated naming
+  os.environ['NCCL_IB_DISABLE'] = '0'  # Enable IB
+  os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'  # Fallback to eth0
+  os.environ['NCCL_IB_GID_INDEX'] = '3'  # Use RoCE v2 (GID index 3)
+  os.environ['NCCL_IB_TC'] = '106'  # Traffic class for RoCE
+  os.environ['NCCL_IB_TIMEOUT'] = '22'  # IB timeout (higher = more tolerant)
+  os.environ['NCCL_IB_RETRY_CNT'] = '7'  # IB retry count
+  os.environ['NCCL_NET_GDR_LEVEL'] = '5'  # Enable GPUDirect RDMA
   os.environ['NCCL_P2P_DISABLE'] = '0'  # Enable P2P
   os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1'  # Better compatibility
-  os.environ['NCCL_LL_THRESHOLD'] = '0'  # Disable LL for compatibility
-  os.environ['NCCL_ALGO'] = 'Ring'  # Use ring algorithm
-  os.environ['NCCL_PROTO'] = 'Simple'  # Use simple protocol
-  os.environ['NCCL_SOCKET_FAMILY'] = 'AF_INET'  # Force IPv4 sockets
 
   # Now import torch after NCCL config is set
   import torch
@@ -136,8 +137,9 @@ def run_training_node_func(rank, world_size):
   print(f"  MASTER_PORT: {MASTER_PORT}")
   print(f"  RANK: {rank}")
   print(f"  WORLD_SIZE: {world_size}")
-  print(f"  NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')}")
-  print(f"  NCCL_SOCKET_IFNAME: {os.environ.get('NCCL_SOCKET_IFNAME')}")
+  print(f"  NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (IB enabled)")
+  print(f"  NCCL_IB_GID_INDEX: {os.environ.get('NCCL_IB_GID_INDEX')}")
+  print(f"  NCCL_IB_TIMEOUT: {os.environ.get('NCCL_IB_TIMEOUT')}")
 
   # STEP 3.5: Pre-flight checks
   print(f"[{rank}] Running pre-flight checks...")