diff --git a/Docker/miniforge_install.sh b/Docker/miniforge_install.sh
index e564af0c..b2ba8c0a 100644
--- a/Docker/miniforge_install.sh
+++ b/Docker/miniforge_install.sh
@@ -60,17 +60,16 @@ source "$HOME/.bashrc"
 "$HOME/conda/bin/conda" install gcc=11 gxx ffmpeg uv cmake make unzip $SYSROOT_PKG "libstdcxx-ng>=11" -y
 
 if [ "$CUDA_VERSION" = "12.8" ]; then
-    "$HOME/conda/bin/uv" pip install torch torchcodec --no-cache-dir --index-url https://download.pytorch.org/whl/cu128
+    "$HOME/conda/bin/uv" pip install ".[cu128]" --no-cache-dir --python "$(which python)"
     "$HOME/conda/bin/conda" install cuda-nvcc=12.8 -y
 elif [ "$CUDA_VERSION" = "12.6" ]; then
-    "$HOME/conda/bin/uv" pip install torch torchcodec --no-cache-dir --index-url https://download.pytorch.org/whl/cu126
+    "$HOME/conda/bin/uv" pip install ".[cu126]" --no-cache-dir --python "$(which python)"
     "$HOME/conda/bin/conda" install cuda-nvcc=12.6 -y
 fi
 
 export PATH="$HOME/conda/bin:$PATH"
 
-"$HOME/conda/bin/uv" pip install psutil ninja packaging wheel "setuptools>=42" einops
-"$HOME/conda/bin/uv" pip install flash-attn -i https://xxxxrt666.github.io/PIP-Index/ --no-build-isolation
+"$HOME/conda/bin/uv" pip install ".[flash-attn]" --python "$(which python)"
 "$HOME/conda/bin/uv" cache clean
 
 rm $LOG_PATH