switch to pytorch docker runtime

Author: zitongzhan

.github/workflows/test-graph-jacobian.yml

@@ -5,28 +5,28 @@ on:
   pull_request:
   workflow_dispatch:
 
+permissions:
+  contents: read
+  packages: read
+
 jobs:
   pytest:
     runs-on: ubuntu-latest
+    container:
+      image: ghcr.io/pytorch/pytorch:latest
     timeout-minutes: 30
     steps:
       - name: Checkout
         uses: actions/checkout@v4
 
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.11"
-          cache: pip
-
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install "torch>=2.1"
           python -m pip install pytest
 
       - name: Run tests
         env:
           PYTHONPATH: ${{ github.workspace }}
         run: |
+          python -c "import torch; print('torch', torch.__version__)"
           python -m pytest -q tests/autograd/test_graph_jacobian.py