Update for new torch

NVIDIA-Merlin · Jun 27, 2024 · 3249d46 · 3249d46
1 parent d8b7dad
commit 3249d46
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/docker/dockerfile.merlin.ctr b/docker/dockerfile.merlin.ctr
@@ -261,8 +261,7 @@ COPY --chown=1000:1000 --from=triton /opt/tritonserver/backends/tensorflow backe
 # Torch Metrics and Lightning (without torch)
 RUN pip install --no-cache-dir tensorflow==2.15.0 wrapt==1.14.0 transformers==4.36.0 \
         && pip uninstall tensorflow keras -y \
-        && pip install --no-cache-dir --no-deps torch torchmetrics \
-        && rm -rf /usr/local/lib/python${PYTHON_VERSION}/dist-packages/torch \
+        && pip install --no-cache-dir --no-deps torchmetrics \
         && rm -rf /usr/local/lib/python${PYTHON_VERSION}/dist-packages/caffe2
 
 # DLFW Tensorflow packages
@@ -285,6 +284,7 @@ RUN ln -s /opt/tritonserver/backends/pytorch/* /usr/local/lib/
 # DLFW Python packages
 COPY --chown=1000:1000 --from=torch /usr/local/lib/python${PYTHON_VERSION}/dist-packages/torch /usr/local/lib/python${PYTHON_VERSION}/dist-packages/torch
 COPY --chown=1000:1000 --from=torch /usr/local/lib/python${PYTHON_VERSION}/dist-packages/torch-*.dist-info /usr/local/lib/python${PYTHON_VERSION}/dist-packages/torch.dist-info/
+COPY --chown=1000:1000 --from=torch /usr/local/lib/python${PYTHON_VERSION}/dist-packages/torchgen /usr/local/lib/python${PYTHON_VERSION}/dist-packages/torchgen
 COPY --chown=1000:1000 --from=torch /usr/local/lib/python${PYTHON_VERSION}/dist-packages/tensorrt /usr/local/lib/python${PYTHON_VERSION}/dist-packages/tensorrt
 COPY --chown=1000:1000 --from=torch /usr/local/lib/python${PYTHON_VERSION}/dist-packages/tensorrt-*.dist-info /usr/local/lib/python${PYTHON_VERSION}/dist-packages/tensorrt.dist-info/