Update Dockerfile

Dockerfile

@@ -67,7 +67,7 @@ RUN ln -sf /usr/bin/python3.10 /usr/bin/python3 && \
 
 # ---------------- PyTorch cu128 (pinned) ----------------
 RUN pip install --index-url https://download.pytorch.org/whl/cu128 \
-    torch==2.8.0+cu128 torchvision==0.23.0+cu128 torchaudio==2.8.0+cu128
+    torch>=2.8.0+cu128 torchvision>=0.23.0+cu128 torchaudio>=2.8.0+cu128
 
 # ---------------- Toolchain, Triton, FA2 (no bnb build) ----------------
 RUN pip install packaging ninja cmake pybind11 scikit-build cython hf_transfer "numpy>=1.24.4"
@@ -76,6 +76,7 @@ RUN pip install packaging ninja cmake pybind11 scikit-build cython hf_transfer "
 RUN pip uninstall -y triton || true && \
     pip install -v --no-build-isolation triton==3.4.0
 
+
 # FlashAttention 2.8.x
 RUN pip install flash-attn==2.8.3 --no-build-isolation || \
     pip install flash-attn==2.8.2 --no-build-isolation || \
@@ -87,7 +88,7 @@ COPY requirements.txt ./requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt
 
 # Pin bnb to avoid surprise CUDA/PTX mismatches (adjust as needed)
-RUN pip install --upgrade "bitsandbytes==0.44.1" || pip install --upgrade "bitsandbytes==0.43.3"
+RUN pip install --upgrade bitsandbytes
 
 # Custom .whl (Apex + dropout_layer_norm)
 RUN echo "Installing custom wheels..." && \