Merge remote-tracking branch 'origin/main' into tensordict-revamp

pytorch · Nov 20, 2023 · bf436e3 · bf436e3
2 parents 53d3e82 + 2fec80e
commit bf436e3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 14 deletions.
diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
@@ -49,16 +49,13 @@ jobs:
 
   benchmark_gpu:
     name: GPU Pytest benchmark
-    runs-on: ubuntu-20.04
-    strategy:
-      matrix:
-        include:
-         - os: linux.4xlarge.nvidia.gpu
-           python-version: 3.8
+    runs-on: linux.g5.4xlarge.nvidia.gpu
     defaults:
       run:
         shell: bash -l {0}
-    container: nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
+    container:
+      image: nvidia/cuda:12.3.0-base-ubuntu22.04
+      options: --gpus all
     steps:
     - name: Install deps
       run: |
@@ -97,8 +94,14 @@ jobs:
         echo /usr/local/bin >> $GITHUB_PATH
     - name: Setup Environment
       run: |
+        python -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
         python -m pip install -e .
         python -m pip install pytest pytest-benchmark
+    - name: check GPU presence
+      run: |
+        python -c """import torch
+        assert torch.cuda.device_count()
+        """
     - name: Run benchmarks
       run: |
         cd benchmarks/

diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
@@ -60,16 +60,13 @@ jobs:
 
   benchmark_gpu:
     name: GPU Pytest benchmark
-    runs-on: ubuntu-20.04
-    strategy:
-      matrix:
-        include:
-         - os: linux.4xlarge.nvidia.gpu
-           python-version: 3.8
+    runs-on: linux.g5.4xlarge.nvidia.gpu
     defaults:
       run:
         shell: bash -l {0}
-    container: nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
+    container:
+      image: nvidia/cuda:12.3.0-base-ubuntu22.04
+      options: --gpus all
     steps:
       - name: Who triggered this?
         run: |
@@ -113,8 +110,14 @@ jobs:
           echo /usr/local/bin >> $GITHUB_PATH
       - name: Setup Environment
         run: |
+          python -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
           python -m pip install -e .
           python -m pip install pytest pytest-benchmark
+      - name: check GPU presence
+        run: |
+          python -c """import torch
+          assert torch.cuda.device_count()
+          """
       - name: Setup benchmarks
         run: |
           echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV