diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py
index 3f2f252..6907499 100644
--- a/dags/test-train-pytorch.py
+++ b/dags/test-train-pytorch.py
@@ -113,18 +113,17 @@ def run_training_node_func(rank, world_size):
   os.environ['WORLD_SIZE'] = str(world_size)
   os.environ['RANK'] = str(rank)
 
-  # NCCL Configuration for InfiniBand
+  # NCCL Configuration - Let NCCL auto-detect best transport
   os.environ['NCCL_DEBUG'] = 'INFO'
   os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
   os.environ['NCCL_BLOCKING_WAIT'] = '1'
-  os.environ['NCCL_IB_DISABLE'] = '0'  # Enable IB
-  os.environ['NCCL_IB_HCA'] = '^mlx5_[0-9]+$'  # Regex for Mellanox HCAs
-  os.environ['NCCL_SOCKET_IFNAME'] = '^ibp[0-9]+s[0-9]+$'  # Regex for IB interfaces
+  # Don't disable IB - let NCCL auto-detect
+  # Don't set SOCKET_IFNAME - let NCCL find the best interface
   os.environ['NCCL_P2P_DISABLE'] = '0'  # Enable P2P
   os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1'  # Better compatibility
-  os.environ['NCCL_NET_GDR_LEVEL'] = '5'  # Enable GPUDirect RDMA
-  os.environ['NCCL_IB_GID_INDEX'] = '3'  # Use RoCE v2 (common for modern IB)
-  os.environ['NCCL_IB_TC'] = '160'  # Traffic class for low latency
+  os.environ['NCCL_IB_GID_INDEX'] = '3'  # Use RoCE v2
+  os.environ['NCCL_IB_TC'] = '106'  # Traffic class (lower value for compatibility)
+  os.environ['NCCL_IB_TIMEOUT'] = '22'  # IB timeout in microseconds
 
   print(f"[{rank}] Environment configured:")
   print(f"  MASTER_ADDR: {MASTER_ADDR}")