diff --git a/dev-requirements.txt b/dev-requirements.txt
index dc8b94831..3b982480e 100644
--- a/dev-requirements.txt
+++ b/dev-requirements.txt
@@ -19,11 +19,11 @@ pytorch-lightning==2.5.0
 tensorboard==2.14.0
 sagemaker==2.230.0
 torch-model-archiver>=0.4.2
-torch>=2.7.0
+torch
 torchmetrics==1.6.3
 torchserve>=0.10.0
-torchtext==0.18.0
-torchvision==0.23.0
+torchtext
+torchvision
 typing-extensions
 ts==0.5.1
 wheel
diff --git a/torchx/runtime/container/Dockerfile b/torchx/runtime/container/Dockerfile
index 004369f96..0a017fcd6 100644
--- a/torchx/runtime/container/Dockerfile
+++ b/torchx/runtime/container/Dockerfile
@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime
+FROM pytorch/pytorch:2.9.0-cuda12.8-cudnn9-runtime
 
 WORKDIR /app