fix autoawq

huggingface · Sep 19, 2024 · 6abc2ba · 6abc2ba
1 parent 073e7dc
commit 6abc2ba
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/docker/cuda-ort/Dockerfile b/docker/cuda-ort/Dockerfile
@@ -30,6 +30,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 # Install PyTorch
 ARG TORCH_CUDA=cu118
 ARG TORCH_VERSION=stable
+ARG IMAGE_FLAVOR
+ENV IMAGE_FLAVOR=${IMAGE_FLAVOR}
 
 RUN if [ "${TORCH_VERSION}" = "stable" ]; then \
     pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \