update PyTorch version to enable torch._scaled_mm rowwise

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
EmbeddedLLM · Jan 28, 2025 · 30f0ecd · 30f0ecd
1 parent 6dc4604
commit 30f0ecd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/Dockerfile.rocm_base b/Dockerfile.rocm_base
@@ -6,7 +6,7 @@ ARG RCCL_BRANCH="648a58d"
 ARG RCCL_REPO="https://github.com/ROCm/rccl"
 ARG TRITON_BRANCH="e5be006"
 ARG TRITON_REPO="https://github.com/triton-lang/triton.git"
-ARG PYTORCH_BRANCH="8d4926e"
+ARG PYTORCH_BRANCH="3a585126"
 ARG PYTORCH_VISION_BRANCH="v0.19.1"
 ARG PYTORCH_REPO="https://github.com/pytorch/pytorch.git"
 ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"