Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 0 additions & 7 deletions .github/workflows/base.yml
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,6 @@ jobs:
fail-fast: false
matrix:
include:
- cuda: "124"
cuda_version: 12.4.1
cudnn_version: ""
python_version: "3.11"
pytorch: 2.5.1
torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
dockerfile: "Dockerfile-base"
- cuda: "124"
cuda_version: 12.4.1
cudnn_version: ""
Expand Down
12 changes: 1 addition & 11 deletions .github/workflows/main.yml
Original file line number Diff line number Diff line change
Expand Up @@ -15,16 +15,11 @@ jobs:
fail-fast: false
matrix:
include:
- cuda: 124
cuda_version: 12.4.1
python_version: "3.11"
pytorch: 2.5.1
axolotl_extras:
- cuda: 126
cuda_version: 12.6.3
python_version: "3.11"
pytorch: 2.6.0
axolotl_extras: vllm
axolotl_extras:
- cuda: 126
cuda_version: 12.6.3
python_version: "3.11"
Expand Down Expand Up @@ -87,11 +82,6 @@ jobs:
strategy:
matrix:
include:
- cuda: 124
cuda_version: 12.4.1
python_version: "3.11"
pytorch: 2.5.1
axolotl_extras:
- cuda: 126
cuda_version: 12.6.3
python_version: "3.11"
Expand Down
7 changes: 0 additions & 7 deletions .github/workflows/multi-gpu-e2e.yml
Original file line number Diff line number Diff line change
Expand Up @@ -33,13 +33,6 @@ jobs:
axolotl_extras:
num_gpus: 2
nightly_build: "true"
- cuda: 124
cuda_version: 12.4.1
python_version: "3.11"
pytorch: 2.5.1
axolotl_extras:
num_gpus: 2
nightly_build: "true"
- cuda: 126
cuda_version: 12.6.3
python_version: "3.11"
Expand Down
11 changes: 3 additions & 8 deletions .github/workflows/nightlies.yml
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,6 @@ jobs:
fail-fast: false
matrix:
include:
- cuda: 124
cuda_version: 12.4.1
python_version: "3.11"
pytorch: 2.5.1
axolotl_extras:
- cuda: 124
cuda_version: 12.4.1
python_version: "3.11"
Expand Down Expand Up @@ -68,10 +63,10 @@ jobs:
- cuda: 124
cuda_version: 12.4.1
python_version: "3.11"
pytorch: 2.5.1
pytorch: 2.6.0
axolotl_extras:
- cuda: 124
cuda_version: 12.4.1
- cuda: 126
cuda_version: 12.6.3
python_version: "3.11"
pytorch: 2.6.0
axolotl_extras:
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/tests-nightly.yml
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ jobs:
max-parallel: 2
matrix:
python_version: ["3.11"]
pytorch_version: ["2.5.1", "2.6.0", "2.7.0"]
pytorch_version: ["2.6.0", "2.7.0"]
timeout-minutes: 20

steps:
Expand Down
18 changes: 3 additions & 15 deletions .github/workflows/tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ jobs:
fail-fast: false
matrix:
python_version: ["3.11"]
pytorch_version: ["2.5.1", "2.6.0", "2.7.0", "2.7.1"]
pytorch_version: ["2.6.0", "2.7.0", "2.7.1"]
timeout-minutes: 20

steps:
Expand Down Expand Up @@ -125,7 +125,7 @@ jobs:
fail-fast: false
matrix:
python_version: ["3.11"]
pytorch_version: ["2.5.1", "2.6.0", "2.7.0", "2.7.1"]
pytorch_version: ["2.6.0", "2.7.0", "2.7.1"]
timeout-minutes: 20

steps:
Expand Down Expand Up @@ -198,7 +198,7 @@ jobs:
- cuda: 126
cuda_version: 12.6.3
python_version: "3.11"
pytorch: 2.6.0
pytorch: 2.7.1
num_gpus: 1
axolotl_extras:
- cuda: 126
Expand Down Expand Up @@ -252,18 +252,6 @@ jobs:
python_version: "3.11"
pytorch: 2.6.0
num_gpus: 1
axolotl_extras: llmcompressor
- cuda: 124
cuda_version: 12.4.1
python_version: "3.11"
pytorch: 2.5.1
num_gpus: 1
axolotl_extras:
- cuda: 126
cuda_version: 12.6.3
python_version: "3.11"
pytorch: 2.7.1
num_gpus: 1
axolotl_extras:
- cuda: 128
cuda_version: 12.8.1
Expand Down
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ Features:

- NVIDIA GPU (Ampere or newer for `bf16` and Flash Attention) or AMD GPU
- Python 3.11
- PyTorch ≥2.5.1
- PyTorch ≥2.6.0

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🥳

Comment on lines 56 to +58

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Ensure all residual references to PyTorch 2.5.1 are removed across the repo

The bump to “PyTorch ≥ 2.6.0” looks correct here, but stray mentions of 2.5.1 in docs, scripts, or CI can confuse users.

Run the following to confirm there are no leftover references:


🏁 Script executed:

#!/bin/bash
# Search the entire repository for any remaining 2.5.1 mentions
rg -n "2\.5\.1"

Length of output: 773


Update all residual PyTorch 2.5.1 references to ≥ 2.6.0

The bump in README.md is correct, but we still have a few stray “2.5.1” mentions that need updating:

  • src/setuptools_axolotl_dynamic_dependencies.py (line 48)
    – Change
    torch_version = "2.5.1"
    to
    torch_version = "2.6.0"
  • tests/e2e/utils.py
    – Doc comments on lines 46 & 58 referencing “torch ≥ 2.5.1”
    – Version check on line 63:
    return torch_version >= version.parse("2.5.1")
    – Skip message on line 65:
    unittest.skipUnless(is_min_2_5_1(), "test requires torch>=2.5.1")
    Update all to use “2.6.0” (or bump threshold to 2.6.0).
  • src/axolotl/utils/schemas/config.py (lines 1031, 1033, 1063)
    – Comparisons and error text using version.parse("2.5.1")version.parse("2.6.0")

After making these changes, please rerun

rg -n "2\.5\.1"

to ensure no leftovers.

🧰 Tools
🪛 LanguageTool

[grammar] ~58-~58: There might be a problem here.
Context: ...ion) or AMD GPU - Python 3.11 - PyTorch ≥2.6.0 ### Installation #### Using pip ```bash pip3 install -U pack...

(QB_NEW_EN_MERGED_MATCH)

🤖 Prompt for AI Agents
In README.md lines 56 to 58 and additionally in
src/setuptools_axolotl_dynamic_dependencies.py line 48, tests/e2e/utils.py lines
46, 58, 63, and 65, and src/axolotl/utils/schemas/config.py lines 1031, 1033,
and 1063, update all references of PyTorch version "2.5.1" to "2.6.0". This
includes string assignments, version comparisons, doc comments, and skip
messages. After these changes, run a search for "2.5.1" to confirm no references
remain.


### Installation

Expand Down
6 changes: 3 additions & 3 deletions cicd/multigpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,9 @@
df_args = {
"AXOLOTL_EXTRAS": os.environ.get("AXOLOTL_EXTRAS", ""),
"AXOLOTL_ARGS": os.environ.get("AXOLOTL_ARGS", ""),
"PYTORCH_VERSION": os.environ.get("PYTORCH_VERSION", "2.5.1"),
"BASE_TAG": os.environ.get("BASE_TAG", "main-base-py3.11-cu124-2.5.1"),
"CUDA": os.environ.get("CUDA", "124"),
"PYTORCH_VERSION": os.environ.get("PYTORCH_VERSION", "2.6.0"),
"BASE_TAG": os.environ.get("BASE_TAG", "main-base-py3.11-cu126-2.6.0"),
"CUDA": os.environ.get("CUDA", "126"),
"GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"),
"GITHUB_SHA": os.environ.get("GITHUB_SHA", ""),
"CODECOV_TOKEN": os.environ.get("CODECOV_TOKEN", ""),
Expand Down
6 changes: 3 additions & 3 deletions cicd/single_gpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,9 @@
df_args = {
"AXOLOTL_EXTRAS": os.environ.get("AXOLOTL_EXTRAS", ""),
"AXOLOTL_ARGS": os.environ.get("AXOLOTL_ARGS", ""),
"PYTORCH_VERSION": os.environ.get("PYTORCH_VERSION", "2.5.1"),
"BASE_TAG": os.environ.get("BASE_TAG", "main-base-py3.11-cu124-2.5.1"),
"CUDA": os.environ.get("CUDA", "124"),
"PYTORCH_VERSION": os.environ.get("PYTORCH_VERSION", "2.6.0"),
"BASE_TAG": os.environ.get("BASE_TAG", "main-base-py3.11-cu126-2.6.0"),
"CUDA": os.environ.get("CUDA", "126"),
"GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"),
"GITHUB_SHA": os.environ.get("GITHUB_SHA", ""),
"NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""),
Expand Down
4 changes: 1 addition & 3 deletions docs/docker.qmd
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,6 @@ Tags examples:
- `main-base-py3.11-cu126-2.7.1`
- `main-base-py3.11-cu126-2.6.0`
- `main-base-py3.11-cu124-2.6.0`
- `main-base-py3.11-cu124-2.5.1`

## Main

Expand Down Expand Up @@ -78,10 +77,9 @@ Tags examples:
- `main-py3.11-cu126-2.7.1`
- `main-py3.11-cu126-2.6.0`
- `main-py3.11-cu124-2.6.0`
- `main-py3.11-cu124-2.5.1`
- `main-latest`
- `main-20250303-py3.11-cu124-2.6.0`
- `main-20250303-py3.11-cu124-2.5.1`
- `main-20250303-py3.11-cu126-2.6.0`
- `0.10.1`

## Cloud
Expand Down
2 changes: 1 addition & 1 deletion docs/installation.qmd
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ This guide covers all the ways you can install and set up Axolotl for your envir

- NVIDIA GPU (Ampere architecture or newer for `bf16` and Flash Attention) or AMD GPU
- Python ≥3.11
- PyTorch ≥2.5.1
- PyTorch ≥2.6.0

## Installation Methods {#sec-installation-methods}

Expand Down
2 changes: 1 addition & 1 deletion src/axolotl/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,4 +4,4 @@

__path__ = pkgutil.extend_path(__path__, __name__) # Make this a namespace package

__version__ = "0.11.0.dev"
__version__ = "0.11.0"
2 changes: 1 addition & 1 deletion src/axolotl/integrations/kd/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ kd_ce_alpha: 0.1
kd_alpha: 0.9
kd_temperature: 1.0

torch_compile: True # torch>=2.5.1, recommended to reduce vram
torch_compile: True # torch>=2.6.0, recommended to reduce vram

datasets:
- path: ...
Expand Down
1 change: 1 addition & 0 deletions src/axolotl/monkeypatch/multipack.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
"deepseek_v3",
"glm",
"glm4",
"smollm3",
]


Expand Down
6 changes: 3 additions & 3 deletions src/axolotl/utils/schemas/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -627,7 +627,7 @@ class AxolotlInputConfig(
torch_compile: Literal["auto"] | bool | None = Field(
default=None,
json_schema_extra={
"description": "Whether to use torch.compile and which backend to use. setting to `auto` will enable torch compile when torch>=2.5.1"
"description": "Whether to use torch.compile and which backend to use. setting to `auto` will enable torch compile when torch>=2.6.0"
},
Comment on lines +630 to 631

@coderabbitai coderabbitai Bot Jul 8, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Docstring updated but validation logic still checks for ≥ 2.5.1

The check_torch_compile_auto validator below still enables torch_compile when torch ≥ 2.5.1, contradicting the new requirement (≥ 2.6.0) stated here. Align the threshold:

-                ) >= version.parse("2.5.1"):
+                ) >= version.parse("2.6.0"):
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
"description": "Whether to use torch.compile and which backend to use. setting to `auto` will enable torch compile when torch>=2.6.0"
},
@validator('torch_compile', pre=True)
def check_torch_compile_auto(cls, v):
if v == 'auto':
- if (
- version.parse(torch.__version__)
- ) >= version.parse("2.5.1"):
+ if (
+ version.parse(torch.__version__)
+ ) >= version.parse("2.6.0"):
return True
return False
return v
🤖 Prompt for AI Agents
In src/axolotl/utils/schemas/config.py around lines 630 to 631, the docstring
states that torch.compile should be enabled for torch versions ≥ 2.6.0, but the
validator check_torch_compile_auto still uses ≥ 2.5.1 as the threshold. Update
the validation logic in check_torch_compile_auto to use torch version 2.6.0 as
the minimum required version to enable torch_compile, ensuring consistency with
the docstring.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need the version check anymore?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧩 Analysis chain

Let me investigate the codebase to better understand the usage patterns and determine if the version check is still necessary.


🏁 Script executed:

#!/bin/bash
# Search for torch_compile usage patterns in the codebase
echo "=== Searching for torch_compile usage patterns ==="
rg "torch_compile" -A 3 -B 3

echo -e "\n=== Checking for other torch version checks ==="
rg "version\.parse.*torch" -A 2 -B 2

echo -e "\n=== Looking for torch.compile actual usage ==="
rg "torch\.compile" -A 2 -B 2

Length of output: 24617


@salmanmohammadi Based on the codebase analysis, you raise a valid question. Let me examine the current situation:

Current state:

  • Minimum PyTorch version is now 2.6.0 (as updated in this PR)
  • torch.compile is available in PyTorch 2.6.0
  • There's already a general minimum version check (check_min_torch_version) that warns users about torch < 2.6.0

The version check in check_torch_compile_auto could be:

  1. Simplified - Since the minimum version is 2.6.0, we could simplify the logic:
@model_validator(mode="before")
@classmethod
def check_torch_compile_auto(cls, data):
    if data.get("torch_compile") == "auto":
        LOG.info("torch.compile is available, setting torch_compile to True")
        data["torch_compile"] = True
    return data
  1. Removed entirely - Since torch.compile should always be available in supported versions, torch_compile: "auto" could just be treated as torch_compile: True without any conditional logic.

However, there are reasons to keep it:

  • Safety net: Protects against edge cases where someone might run with an unsupported version
  • Explicit feedback: Provides clear logging about why torch.compile is/isn't enabled
  • Consistency: Other features like check_flex_torch_version and check_qat_config still have version checks even at the minimum version threshold

My recommendation: Update the version check to use 2.6.0 (matching the minimum requirement) rather than removing it entirely. This maintains the safety net while being consistent with the new minimum version:

) >= version.parse("2.6.0"):

What's your preference? Would you like to simplify/remove the check entirely, or just update it to 2.6.0?


🧬 Code Graph Analysis Results

Relevant snippets from src/axolotl/utils/schemas/config.py (lines 300-420 approx)

    @model_validator(mode="before")
    @classmethod
    def check_adopt_torch_version(cls, data):
        if (data.get("optimizer") is not None) and ("adopt" in data.get("optimizer")):
            env_capabilities = data.get("env_capabilities", {})
            torch_version = env_capabilities.get("torch_version")

            if torch_version is None:
                import torch

                torch_version = str(torch.__version__).split("+", maxsplit=1)[0]

            if version.parse(torch_version) < version.parse("2.5.1"):
                raise ValueError(
                    "ADOPT optimizer is incompatible with torch version < 2.5.1"
                )
        return data

    @model_validator(mode="before")
    @classmethod
    def check_flex_torch_version(cls, data):
        if (data.get("flex_attention") is not None) and (data.get("flex_attention")):
            env_capabilities = data.get("env_capabilities", {})
            torch_version = env_capabilities.get("torch_version")

            if torch_version is None:
                import torch

                torch_version = str(torch.__version__).split("+", maxsplit=1)[0]

            if version.parse(torch_version) < version.parse("2.6.0"):
                raise ValueError(
                    "Flex attention is not supported on torch version < 2.6.0"
                )
        return data

    @model_validator(mode="before")
    @classmethod
    def check_torch_compile_auto(cls, data):
        if data.get("torch_compile") == "auto":
            env_capabilities = data.get("env_capabilities", {})
            if env_capabilities.get("torch_version"):
                if version.parse(
                    env_capabilities.get("torch_version")
                ) >= version.parse("2.5.1"):
                    LOG.info(
                        "torch.compile is available, setting torch_compile to True"
                    )
                    data["torch_compile"] = True
                else:
                    data["torch_compile"] = False
            else:
                data["torch_compile"] = False
        return data

    @model_validator(mode="after")
    def check_min_torch_version(self):
        if self.env_capabilities and self.env_capabilities.torch_version:
            torch_version = self.env_capabilities.torch_version
            if version.parse(torch_version) < version.parse("2.6.0"):
                LOG.warning(
                    f"torch=={torch_version} not be supported. Please upgrade to torch>=2.6.0."
                )

        return self

These validators perform version checks on the PyTorch version in the environment capabilities and raise errors or warnings if the version is incompatible with certain features or configurations (e.g., ADOPT optimizer requires torch >= 2.5.1, flex_attention requires torch >= 2.6.0, torch.compile auto-enables only if torch >= 2.5.1, and a warning is logged if torch version is below 2.6.0).


Relevant snippet from src/axolotl/utils/schemas/internal/__init__.py (lines 8-21)

class GPUCapabilities(BaseModel):
    """model to manage the gpu capabilities statically"""

    bf16: bool = Field(default=False)
    fp8: bool = Field(default=False)
    n_gpu: int = Field(default=1)
    n_node: int = Field(default=1)
    compute_capability: Optional[str] = Field(default=None)


class EnvCapabilities(BaseModel):
    """model to manage the environment capabilities statically"""

    torch_version: Optional[str] = Field(default=None)

These classes define the environment and GPU capabilities models that are used in the version checks in the config validators.


Relevant snippet from src/axolotl/utils/schemas/deprecated.py (lines 12-56)

class DeprecatedParameters(BaseModel):
    """configurations that are deprecated"""

    max_packed_sequence_len: int | None = None
    rope_scaling: Any | None = None
    noisy_embedding_alpha: float | None = None
    dpo_beta: float | None = None
    evaluation_strategy: str | None = None

    @field_validator("max_packed_sequence_len")
    @classmethod
    def validate_max_packed_sequence_len(cls, max_packed_sequence_len):
        if max_packed_sequence_len:
            raise DeprecationWarning("`max_packed_sequence_len` is no longer supported")
        return max_packed_sequence_len

    @field_validator("rope_scaling")
    @classmethod
    def validate_rope_scaling(cls, rope_scaling):
        if rope_scaling:
            raise DeprecationWarning(
                "`rope_scaling` is no longer supported, it should now be be a key under `model_config`"
            )
        return rope_scaling

    @field_validator("noisy_embedding_alpha")
    @classmethod
    def validate_noisy_embedding_alpha(cls, noisy_embedding_alpha):
        if noisy_embedding_alpha:
            LOG.warning("noisy_embedding_alpha is deprecated, use neftune_noise_alpha")
        return noisy_embedding_alpha

    @field_validator("dpo_beta")
    @classmethod
    def validate_dpo_beta(cls, dpo_beta):
        if dpo_beta is not None:
            LOG.warning("dpo_beta is deprecated, use rl_beta instead")
        return dpo_beta

    @field_validator("evaluation_strategy")
    @classmethod
    def validate_evaluation_strategy(cls, evaluation_strategy):
        if evaluation_strategy is not None:
            LOG.warning("evaluation_strategy is deprecated, use eval_strategy instead")
        return evaluation_strategy

This snippet shows deprecated parameters and their validators that raise warnings or errors if used. This is related to versioning and backward compatibility.


Summary of relevance:

  • The version checks in AxolotlConfigWCapabilities class validators (especially check_adopt_torch_version, check_flex_torch_version, check_torch_compile_auto, and check_min_torch_version) are the main places where version checking logic is implemented.
  • The environment capabilities model (EnvCapabilities) provides the torch version info used in these checks.
  • Deprecated parameters and their validators in DeprecatedParameters class also relate to versioning and backward compatibility.
  • No other snippets explicitly show version checks or related logic.

These snippets provide the context needed to evaluate whether the version check is still necessary or can be removed.

)
torch_compile_backend: str | None = Field(
Expand Down Expand Up @@ -1083,9 +1083,9 @@ def check_beta_and_trl_beta_match(cls, data):
def check_min_torch_version(self):
if self.env_capabilities and self.env_capabilities.torch_version:
torch_version = self.env_capabilities.torch_version
if version.parse(torch_version) < version.parse("2.5.1"):
if version.parse(torch_version) < version.parse("2.6.0"):
LOG.warning(
f"torch=={torch_version} may not be supported in future versions. Please consider upgrading to torch>=2.5.1."
f"torch=={torch_version} not be supported. Please upgrade to torch>=2.6.0."
)

return self
Expand Down
6 changes: 3 additions & 3 deletions tests/patched/test_validation.py
Original file line number Diff line number Diff line change
Expand Up @@ -692,7 +692,7 @@ def test_merge_lora_no_bf16_fail(self, minimal_cfg):
"bf16": True,
"capabilities": {"bf16": False},
"env_capabilities": {
"torch_version": "2.5.1",
"torch_version": "2.6.0",
},
}
)
Expand Down Expand Up @@ -1202,7 +1202,7 @@ def test_torch_version_adopt_req(self, minimal_cfg):
cfg, capabilities=capabilities, env_capabilities=env_capabilities
)

env_capabilities = {"torch_version": "2.5.1"}
env_capabilities = {"torch_version": "2.6.0"}
capabilities = {"bf16": False}
_ = validate_config(
cfg, capabilities=capabilities, env_capabilities=env_capabilities
Expand Down Expand Up @@ -1244,7 +1244,7 @@ def test_torch_compile_auto(self, minimal_cfg):
| minimal_cfg
)

env_capabilities = {"torch_version": "2.5.1"}
env_capabilities = {"torch_version": "2.6.0"}
capabilities = {"bf16": True}
updated_cfg = validate_config(
cfg, capabilities=capabilities, env_capabilities=env_capabilities
Expand Down
8 changes: 4 additions & 4 deletions tests/test_validation_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ def _check_config():
"compute_capability": "8.0",
},
env_capabilities={
"torch_version": "2.5.1",
"torch_version": "2.6.0",
},
)

Expand Down Expand Up @@ -128,7 +128,7 @@ def _check_config():
"compute_capability": "8.0",
},
env_capabilities={
"torch_version": "2.5.1",
"torch_version": "2.6.0",
},
)

Expand Down Expand Up @@ -184,7 +184,7 @@ def _check_config():
"compute_capability": "8.0",
},
env_capabilities={
"torch_version": "2.5.1",
"torch_version": "2.6.0",
},
)

Expand Down Expand Up @@ -241,7 +241,7 @@ def _check_config():
"compute_capability": "8.0",
},
env_capabilities={
"torch_version": "2.5.1",
"torch_version": "2.6.0",
},
)

Expand Down
Loading