diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py
index 8afc74c..3f2f252 100644
--- a/dags/test-train-pytorch.py
+++ b/dags/test-train-pytorch.py
@@ -113,14 +113,18 @@ def run_training_node_func(rank, world_size):
   os.environ['WORLD_SIZE'] = str(world_size)
   os.environ['RANK'] = str(rank)
 
-  # NCCL Configuration - Start with basic ethernet config (IB is causing issues)
+  # NCCL Configuration for InfiniBand
   os.environ['NCCL_DEBUG'] = 'INFO'
   os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
   os.environ['NCCL_BLOCKING_WAIT'] = '1'
-  os.environ['NCCL_IB_DISABLE'] = '1'  # Disable IB for now - use ethernet
-  os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'  # Use regular ethernet
+  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
   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
 
   print(f"[{rank}] Environment configured:")
   print(f"  MASTER_ADDR: {MASTER_ADDR}")