name: Setup

inputs:

  python-version:

    required: false

    default: '3.9'

  torch-version:

    required: false

    default: '2.0.0'

  cuda-version:

    required: false

    default: cpu

  full_install:

    required: false

    default: true

runs:

  using: composite

  steps:

    - name: Set up Python ${{ inputs.python-version }}

      uses: actions/setup-python@v4.3.0

      with:

        python-version: ${{ inputs.python-version }}

        check-latest: true

        cache: pip

        cache-dependency-path: |

          pyproject.toml

    - name: Install PyTorch ${{ inputs.torch-version }}+${{ inputs.cuda-version }}

      run: |

        pip install torch==${{ inputs.torch-version }} --extra-index-url https://download.pytorch.org/whl/${{ inputs.cuda-version }}

        python -c "import torch; print('PyTorch:', torch.__version__)"

        python -c "import torch; print('CUDA available:', torch.cuda.is_available())"

        python -c "import torch; print('CUDA:', torch.version.cuda)"

      shell: bash

    - name: Install extension packages

      if: ${{ inputs.full_install == 'true' }}

      run: |

        pip install scipy

        pip install --no-index --upgrade torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-${{ inputs.torch-version }}+${{ inputs.cuda-version }}.html

      shell: bash