diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py
index 8393b4a..5b88623 100644
--- a/dags/test-train-pytorch.py
+++ b/dags/test-train-pytorch.py
@@ -40,13 +40,10 @@ def run_training_node_func(rank, world_size):
   import time
 
   # CRITICAL: Set NCCL environment variables BEFORE importing torch
-  # NOTE: IB is disabled due to subnet manager PKEY mismatch (test-2 only has limited membership 0x7fff)
-  # Cluster admin needs to configure opensm partitions to assign full membership PKEYs to both ports
   os.environ['NCCL_DEBUG'] = 'INFO'
   os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
-  os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1'  # Updated naming
-  os.environ['NCCL_IB_DISABLE'] = '1'  # Disable IB - use ethernet (PKEY misconfiguration)
-  os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'  # Use ethernet (working)
+  os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1'
+  os.environ['NCCL_IB_DISABLE'] = '0'  # Enable InfiniBand
   os.environ['NCCL_P2P_DISABLE'] = '0'  # Enable P2P
   os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1'  # Better compatibility
 
@@ -133,8 +130,7 @@ 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')} (ethernet mode)")
-  print(f"  NCCL_SOCKET_IFNAME: {os.environ.get('NCCL_SOCKET_IFNAME')}")
+  print(f"  NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (IB enabled)")
 
   # STEP 3.5: Pre-flight checks
   print(f"[{rank}] Running pre-flight checks...")