From 9c032a47d52a8d6f001339375b924275bb0f1d3d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Fri, 10 Jan 2025 09:29:39 +0100
Subject: [PATCH] build: Fix mamba install
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
---
 Dockerfile.ci                     | 3 ++-
 requirements/requirements_nlp.txt | 1 -
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/Dockerfile.ci b/Dockerfile.ci
index 2b8a55eac987..628f610ab015 100644
--- a/Dockerfile.ci
+++ b/Dockerfile.ci
@@ -35,10 +35,11 @@ WORKDIR /workspace
 
 # Install Mamba Dependancy
 ARG CAUSAL_CONV_TAG=v1.2.2.post1
+ARG MAMBA_TAG=v2.2.2
 
 RUN <<"EOF" bash -ex
 # Mamba dependancy installation
-
+pip install git+https://github.com/state-spaces/mamba.git@${MAMBA_TAG}
 git clone --depth 1 --branch ${CAUSAL_CONV_TAG} https://github.com/Dao-AILab/causal-conv1d && \
   cd causal-conv1d && \
   python setup.py install && \
diff --git a/requirements/requirements_nlp.txt b/requirements/requirements_nlp.txt
index 6a86dacbfefb..cd30431e8ca9 100644
--- a/requirements/requirements_nlp.txt
+++ b/requirements/requirements_nlp.txt
@@ -8,7 +8,6 @@ gdown
 h5py
 ijson
 jieba
-mamba-ssm==2.2.2; sys_platform == 'linux'
 markdown2
 matplotlib>=3.3.2
 #megatron_core>0.6.0 # add back once mcore on pypi is compatible again