From 3a8cea0653846f91e8519ca37484e881238213e2 Mon Sep 17 00:00:00 2001
From: Eitan Turok <150733043+eitanturok@users.noreply.github.com>
Date: Mon, 19 Aug 2024 17:11:17 -0400
Subject: [PATCH] test torch 2.4

---
 .github/workflows/pr-gpu.yaml | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/.github/workflows/pr-gpu.yaml b/.github/workflows/pr-gpu.yaml
index 1ca8d5b..38433e6 100644
--- a/.github/workflows/pr-gpu.yaml
+++ b/.github/workflows/pr-gpu.yaml
@@ -29,6 +29,14 @@ jobs:
           gpu_num: 2
           python_version: 3.11
           container: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
+        - name: "python3.11-pytorch2.4.0-gpus1"
+          gpu_num: 1
+          python_version: 3.11
+          container: mosaicml/pytorch:2.4.0_cu121-python3.11-ubuntu20.04
+        - name: "python3.11-pytorch2.4.0-gpus2"
+          gpu_num: 2
+          python_version: 3.11
+          container: mosaicml/pytorch:2.4.0_cu121-python3.11-ubuntu20.04
     steps:
     - name: Run PR GPU tests
       uses: mosaicml/ci-testing/.github/actions/pytest-gpu@v0.1.2