From 2fa214ba74406ee4ec96c5b9c80773c3085b4636 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Mon, 2 Dec 2024 17:44:00 -0500 Subject: [PATCH] build causal_conv1d and mamba-ssm into the base image --- docker/Dockerfile-base | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base index 3f13bba30a..7eab3b3e43 100644 --- a/docker/Dockerfile-base +++ b/docker/Dockerfile-base @@ -29,7 +29,9 @@ ENV PATH="/root/miniconda3/envs/py${PYTHON_VERSION}/bin:${PATH}" WORKDIR /workspace RUN python3 -m pip install --upgrade pip && pip3 install packaging && \ - python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} --extra-index-url https://download.pytorch.org/whl/cu$CUDA + python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} --extra-index-url https://download.pytorch.org/whl/cu$CUDA && \ + python3 -m pip install --no-cache-dir "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" && \ + python3 -m pip install --no-cache-dir "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" RUN git lfs install --skip-repo && \ pip3 install awscli && \