fix: Updating Dockerfile to pin versions and fix the exampel (#505)

aws-samples · Nov 23, 2024 · f81f48f · f81f48f
1 parent d911db0
commit f81f48f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/3.test_cases/10.FSDP/Dockerfile b/3.test_cases/10.FSDP/Dockerfile
@@ -3,8 +3,8 @@ FROM public.ecr.aws/hpc-cloud/nccl-tests:latest
 
 RUN apt update && apt install -y nvtop
 
-RUN pip install torchvision torchaudio transformers datasets fsspec==2023.9.2 python-etcd numpy==1.*
-RUN pip install torch==2.2.0+cu121 --index-url https://download.pytorch.org/whl/cu121
+RUN pip install torchvision torchaudio transformers==4.46.1 datasets fsspec==2023.9.2 python-etcd numpy==1.*
+RUN pip install torch==2.5.1+cu121 --index-url https://download.pytorch.org/whl/cu121
 
 RUN mkdir /checkpoints