From 444a86f5dc12edbcb8f11cf345d27f645056deaa Mon Sep 17 00:00:00 2001
From: Jiahao Su <damow890@gmail.com>
Date: Sat, 30 Aug 2025 11:05:45 +0800
Subject: [PATCH 1/2] Update pytorch version

---
 arg.json           | 16 ++++++++--------
 pytorch/Dockerfile | 18 +++++++++---------
 pytorch/README.md  |  8 ++++----
 3 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/arg.json b/arg.json
index 5ecc9d5..b0c8a8d 100644
--- a/arg.json
+++ b/arg.json
@@ -379,24 +379,24 @@
   ],
   "pytorch": [
     {
-      "pytorch_version": "2.1.0",
-      "cann_tag": "8.0.rc1-910b-ubuntu22.04-py3.8",
+      "pytorch_version": "2.6.0",
+      "cann_tag": "8.2.rc1-910b-ubuntu22.04-py3.11",
       "tags": {
         "common": [
-          "2.1.0",
-          "ubuntu-python3.8-cann8.0.rc1-pytorch2.1.0"
+          "2.6.0",
+          "ubuntu-python3.11-cann8.2.rc1-pytorch2.6.0"
         ],
         "ascendhub": [
         ]
       }
     },
     {
-      "pytorch_version": "2.2.0",
-      "cann_tag": "8.0.rc1-910b-ubuntu22.04-py3.8",
+      "pytorch_version": "2.7.1",
+      "cann_tag": "8.2.rc1-910b-ubuntu22.04-py3.11",
       "tags": {
         "common": [
-          "2.2.0",
-          "ubuntu-python3.8-cann8.0.rc1-pytorch2.2.0"
+          "2.7.1",
+          "ubuntu-python3.11-cann8.2.rc1-pytorch2.7.1"
         ],
         "ascendhub": [
         ]
diff --git a/pytorch/Dockerfile b/pytorch/Dockerfile
index a793fe1..2826738 100644
--- a/pytorch/Dockerfile
+++ b/pytorch/Dockerfile
@@ -2,20 +2,20 @@ ARG BASE_VERSION=latest
 FROM ascendai/cann:${BASE_VERSION} as official
 
 # Arguments
-ARG PYTORCH_VERSION=2.2.0
+ARG PYTORCH_VERSION=2.7.1
 
 # Change the default shell
 SHELL [ "/bin/bash", "-c" ]
 
 # Install pytorch, torch-npu and related packages
-RUN if [ "${PYTORCH_VERSION}" == "2.1.0" ]; then \
-        TORCH_VISION_VERSION=0.16.0; \
-        TORCH_AUDIO_VERSION=2.1.0; \
-        TORCH_NPU_VERSION=2.1.0; \
-    elif [ "${PYTORCH_VERSION}" == "2.2.0" ]; then \
-        TORCH_VISION_VERSION=0.17.0; \
-        TORCH_AUDIO_VERSION=2.2.0; \
-        TORCH_NPU_VERSION=2.2.0; \
+RUN if [ "${PYTORCH_VERSION}" == "2.6.0" ]; then \
+        TORCH_VISION_VERSION=0.21.0; \
+        TORCH_AUDIO_VERSION=2.6.0; \
+        TORCH_NPU_VERSION=2.6.0; \
+    elif [ "${PYTORCH_VERSION}" == "2.7.1" ]; then \
+        TORCH_VISION_VERSION=0.22.1; \
+        TORCH_AUDIO_VERSION=2.7.1; \
+        TORCH_NPU_VERSION=2.7.1rc1; \
     else \
         echo "Not supported version: ${PYTORCH_VERSION}."; \
         exit 1; \
diff --git a/pytorch/README.md b/pytorch/README.md
index 059cd8e..9b6e3a6 100644
--- a/pytorch/README.md
+++ b/pytorch/README.md
@@ -12,8 +12,8 @@ The pytorch image is based on the [cann](../cann) image. These packages are inst
 
 Here are available tags and the build args can be found at [arg.json](../arg.json).
 
-- `2.1.0`
-- `2.2.0`
+- `2.6.0`
+- `2.7.1`
 
 > [!NOTE]
 >
@@ -36,7 +36,7 @@ docker run \
     -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
     -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
-    -it ascendai/pytorch:2.2.0 bash
+    -it ascendai/pytorch:2.7.1 bash
 ```
 
 ## Build
@@ -56,6 +56,6 @@ docker build \
     -t ascendai/pytorch:latest \
     -f pytorch/new.Dockerfile \
     --build-arg BASE_VERSION=latest \
-    --build-arg PYTORCH_VERSION=2.2.0 \
+    --build-arg PYTORCH_VERSION=2.7.1 \
     pytorch/
 ```

From 62496e704099f2a5bb688f7be6041d0ad2ae93cb Mon Sep 17 00:00:00 2001
From: Jiahao Su <damow890@gmail.com>
Date: Sat, 30 Aug 2025 17:34:45 +0800
Subject: [PATCH 2/2] update

---
 .github/workflows/docker.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
index 17b1a24..b6d665a 100644
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@@ -122,7 +122,7 @@ jobs:
     name: pytorch
     needs:
       - prepare
-      - cann
+      # - cann
     if: |
       !cancelled() &&
       contains(needs.prepare.outputs.targets, 'pytorch') &&