step: install_torch
description: "安装 PyTorch（含 CUDA 支持验证）"
validation:
  script: |
    python3 -c "import torch; assert torch.cuda.is_available(); print(torch.__version__, torch.cuda.get_device_capability(0))"
  expected_pattern: "({gpu_cc})"

branches:
  - id: pip_nvidia_index
    description: "pip install torch（NVIDIA CUDA index，自动匹配系统 CUDA 版本）"
    command_template: "pip install torch --index-url https://download.pytorch.org/whl/cu{cuda_ver_short}"
    fallback_on: ["driver too old", "CUDA driver version is insufficient"]
    note: |
      - 从 nvidia-smi 自动检测 CUDA 版本（如 12.8 → cu128, 12.4 → cu124）
      - 国内需配合阿里云镜像加速

  - id: pip_nvidia_index_mirror
    description: "pip install torch（NVIDIA CUDA index + 阿里云镜像加速）"
    command_template: "pip install torch --index-url https://download.pytorch.org/whl/cu{cuda_ver_short} -i https://mirrors.aliyun.com/pypi/simple/"
    note: "阿里云镜像 aarch64 wheel 覆盖可能不全"

  - id: pip_default
    description: "pip install torch（默认 PyPI，aarch64 上通常给 CPU-only）"
    command: "pip install torch"
    fallback_on: ["No matching distribution", "CPU-only", "ModuleNotFoundError: No module named 'torch'"]

  - id: conda
    description: "conda install pytorch（自动匹配 CUDA 版本）"
    command_template: "conda install pytorch torchvision torchaudio pytorch-cuda={cuda_ver} -c pytorch -c nvidia"
    fallback_on: ["conda: command not found", "channel not found"]
    note: "aarch64 conda channel 支持有限"