diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
index 03ede05..ef4556e 100644
--- a/.devcontainer/Dockerfile
+++ b/.devcontainer/Dockerfile
@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-devel as compile_server
+FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel as compile_server
 WORKDIR /workspace
 ENV CUDA_HOME /usr/local/cuda
 RUN <<EOF
@@ -10,6 +10,7 @@ apt update -y &&  apt install -y  --no-install-recommends \
     g++ \
     cmake && 
 rm -rf /var/lib/apt/lists/* &&
+pip install --upgrade pip &&
 pip install ninja pyproject numpy cpufeature &&
 pip install flash-attn &&
 cp /usr/lib/x86_64-linux-gnu/libstdc++.so.6 /opt/conda/lib/
diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml
index 60df01f..5eb5cec 100644
--- a/.github/workflows/docker-image.yml
+++ b/.github/workflows/docker-image.yml
@@ -3,6 +3,14 @@ name: DockerHub CI
 on:
   release:
     types: [published]
+  workflow_dispatch:
+    inputs:
+      choose:
+        description: 'Will you push the image to DockerHub? 0 for No, 1 for Yes'
+        required: true
+        default: '0'
+        type: string
+
   # push:
   #   branches:
   #     - main
diff --git a/.gitignore b/.gitignore
index 7e186ef..74a7725 100644
--- a/.gitignore
+++ b/.gitignore
@@ -25,3 +25,4 @@ book
 ktransformers/tests/chat_txt.txt
 mmlu_result*
 ktransformers/ktransformers_ext/cuda_musa/
+test_prompt.txt
diff --git a/Dockerfile b/Dockerfile
index 1807150..473fbcb 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -10,7 +10,7 @@ EOF
 
 
 
-FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-devel as compile_server
+FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel as compile_server
 ARG CPU_INSTRUCT=NATIVE
 WORKDIR /workspace
 ENV CUDA_HOME /usr/local/cuda
@@ -27,6 +27,7 @@ rm -rf /var/lib/apt/lists/* &&
 cd ktransformers &&
 git submodule init &&
 git submodule update &&
+pip install --upgrade pip &&
 pip install ninja pyproject numpy cpufeature &&
 pip install flash-attn &&
 CPU_INSTRUCT=${CPU_INSTRUCT}  KTRANSFORMERS_FORCE_BUILD=TRUE TORCH_CUDA_ARCH_LIST="8.0;8.6;8.7;8.9;9.0+PTX" pip install . --no-build-isolation --verbose &&