Fix DeepSpeed CI (#22194)

* Deal with torch-tensorrt --------- Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

Fix DeepSpeed CI (#22194)
* Deal with torch-tensorrt --------- Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
1c4a9acc · Yih-Dar · GitHub · 7c4999e4 · 1c4a9acc
Unverified Commit 1c4a9acc authored Mar 16, 2023 by Yih-Dar Committed by GitHub Mar 16, 2023
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 3 deletions

docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile +2 -3

No files found.
--- a/docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile
+++ b/docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile
@@ -24,9 +24,8 @@ RUN python3 -m pip install --no-cache-dir -U torch torchvision torchaudio --inde

 RUN python3 -m pip install --no-cache-dir ./transformers[deepspeed-testing]

-# This will uninstall torch 2.0.0
-# TODO: uncomment the following line once `torch-tensorrt` is ready for `torch 2.0.0`
-# RUN python3 -m pip install torch-tensorrt==1.3.0 --find-links https://github.com/pytorch/TensorRT/releases/expanded_assets/v1.3.0
+# Uninstall `torch-tensorrt` shipped with the base image
+RUN python3 -m pip uninstall -y torch-tensorrt

 # recompile apex
 RUN python3 -m pip uninstall -y apex