diff --git a/.github/workflows/pr-gpu.yaml b/.github/workflows/pr-gpu.yaml index dda6e68d..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/** @@ -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..4a8274ad 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; docker image doesn't have triton installed so installing here. 'stanford-stk==0.7.1', ]