From a6d613629660fbcd84eabc12f64e2a3e1712b716 Mon Sep 17 00:00:00 2001
From: Rithwik Ediga Lakhamsani <rithwik.ediga@databricks.com>
Date: Thu, 6 Mar 2025 15:50:15 -0800
Subject: [PATCH 1/3] upgrading torch to 2.6

---
 .github/workflows/pr-gpu.yaml | 8 ++++----
 pyproject.toml                | 2 +-
 setup.py                      | 4 ++--
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/pr-gpu.yaml b/.github/workflows/pr-gpu.yaml
index dda6e68d..b97c2f45 100644
--- a/.github/workflows/pr-gpu.yaml
+++ b/.github/workflows/pr-gpu.yaml
@@ -21,14 +21,14 @@ jobs:
       fail-fast: false
       matrix:
         include:
-        - name: "python3.11-pytorch2.5.1-gpus1"
+        - name: "python3.11-pytorch2.6.0-gpus1"
           gpu_num: 1
           python_version: 3.11
-          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
-        - name: "python3.11-pytorch2.5.1-gpus2"
+          container: mosaicml/pytorch:2.6.0_cu124-python3.11-ubuntu20.04
+        - name: "python3.11-pytorch2.6.0-gpus2"
           gpu_num: 2
           python_version: 3.11
-          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.6.0_cu124-python3.11-ubuntu20.04
     steps:
     - name: Run PR GPU tests
       uses: mosaicml/ci-testing/.github/actions/pytest-gpu@v0.1.2
diff --git a/pyproject.toml b/pyproject.toml
index 28b9135c..7fcc670b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@
 
 # build requirements
 [build-system]
-requires = ["setuptools < 70.0.0", "torch >= 2.5.1, < 2.5.2"]
+requires = ["setuptools < 70.0.0", "torch >= 2.6.0, < 2.6.1"]
 build-backend = "setuptools.build_meta"
 
 # Pytest
diff --git a/setup.py b/setup.py
index 5338d6b3..13a837d4 100644
--- a/setup.py
+++ b/setup.py
@@ -62,8 +62,8 @@
 install_requires = [
     'numpy>=1.21.5,<2.1.0',
     'packaging>=21.3.0,<24.2',
-    'torch>=2.5.1,<2.5.2',
-    'triton>=2.1.0',
+    'torch>=2.6.0,<2.6.1',
+    'triton>=3.2.0,<3.3.0', # torch 2.6.0 requires ==3.2.0
     'stanford-stk==0.7.1',
 ]
 

From 07d642ca461be208d7b829919d71679d5da52b9f Mon Sep 17 00:00:00 2001
From: Rithwik Ediga Lakhamsani <rithwik.ediga@databricks.com>
Date: Thu, 6 Mar 2025 15:54:15 -0800
Subject: [PATCH 2/3] comment change

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 13a837d4..4a8274ad 100644
--- a/setup.py
+++ b/setup.py
@@ -63,7 +63,7 @@
     'numpy>=1.21.5,<2.1.0',
     'packaging>=21.3.0,<24.2',
     'torch>=2.6.0,<2.6.1',
-    'triton>=3.2.0,<3.3.0', # torch 2.6.0 requires ==3.2.0
+    'triton>=3.2.0,<3.3.0', # torch 2.6.0 requires ==3.2.0; docker image doesn't have triton installed so installing here.
     'stanford-stk==0.7.1',
 ]
 

From 15cb71aa15951af0e458b87ac2cf31ec1b40e21f Mon Sep 17 00:00:00 2001
From: Rithwik Ediga Lakhamsani <rithwik.ediga@databricks.com>
Date: Mon, 10 Mar 2025 13:57:26 -0700
Subject: [PATCH 3/3] added pull request target

---
 .github/workflows/pr-gpu.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/pr-gpu.yaml b/.github/workflows/pr-gpu.yaml
index b97c2f45..610d7fa6 100644
--- a/.github/workflows/pr-gpu.yaml
+++ b/.github/workflows/pr-gpu.yaml
@@ -4,7 +4,7 @@ on:
     branches:
     - main
     - release/*
-  pull_request_target:
+  pull_request:
     branches:
     - main
     - release/**