diff --git a/Dockerfile_amd b/Dockerfile_amd
index efc80234..129a58fc 100644
--- a/Dockerfile_amd
+++ b/Dockerfile_amd
@@ -92,14 +92,17 @@ RUN chmod +x ~/mambaforge.sh && \
 # Install flash-attention, torch dependencies
 RUN pip install numpy einops ninja --no-cache-dir
 
-RUN conda install mkl-static mkl-include
+RUN conda install intel::mkl-static intel::mkl-include
 RUN pip uninstall -y triton && \
     git clone --depth 1 --single-branch https://github.com/ROCm/triton.git && \
     cd triton/python && \
     pip install .
 
-RUN git clone --depth 1 --recursive --single-branch --branch main https://github.com/pytorch/pytorch.git pytorch && cd pytorch && pip install -r requirements.txt --no-cache-dir && \
-    git checkout da320214e66b5af0f7db8fd18a64dbb519d17b27
+RUN git clone --depth 1 --recursive --single-branch --branch main https://github.com/pytorch/pytorch.git pytorch && \
+    cd pytorch && git fetch --depth 1 origin da320214e66b5af0f7db8fd18a64dbb519d17b27 && \
+    git checkout da320214e66b5af0f7db8fd18a64dbb519d17b27 && \
+    pip install -r requirements.txt --no-cache-dir
+
 
 ARG _GLIBCXX_USE_CXX11_ABI="1"
 ARG CMAKE_PREFIX_PATH="/opt/conda"