_build_torch_with_cuda.yml

name: build-torch-with-cuda
on:
  workflow_call:
    inputs:
      dev-image:
        required: true
        type: string
        description: Base image for builds
      torch-commit:
        required: true
        type: string
        description: torch-commit
      runner:
        required: false
        type: string
        description: Runner type for the test
        default: linux.12xlarge
jobs:
  build:
    runs-on: ${{ inputs.runner }}
    container:
      image: ${{ inputs.dev-image }}
    env:
      _GLIBCXX_USE_CXX11_ABI: 0
    steps:
      - name: Checkout actions
        uses: actions/checkout@v4
        with:
          sparse-checkout: |
            .github/workflows/setup
          path: .actions
      - name: Setup
        uses: ./.actions/.github/workflows/setup
        with:
          torch-commit: ${{ inputs.torch-commit }}
          cuda: true
      - name: Checkout PyTorch Repo
        uses: actions/checkout@v4
        with:
          repository: pytorch/pytorch
          path: pytorch
          ref: ${{ inputs.torch-commit }}
          submodules: recursive
      - name: Build PyTorch with CUDA enabled
        shell: bash
        run: |
          cd pytorch
          TORCH_CUDA_ARCH_LIST="5.2;8.6" USE_CUDA=1 MAX_JOBS="$(nproc --ignore=4)" python setup.py bdist_wheel
      - name: Upload wheel
        uses: actions/upload-artifact@v4
        with:
          name: torch-with-cuda
          path: pytorch/dist/*.whl