name: Unit-tests on Linux GPU

on:
  pull_request:
  push:
    branches:
      - nightly
      - main
      - release/*
  workflow_dispatch:

env:
  CHANNEL: "nightly"

jobs:
  tests:
    strategy:
      matrix:
        python_version: ["3.8"]
        cuda_arch_version: ["11.6"]
      fail-fast: false
    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
    with:
      runner: linux.g5.4xlarge.nvidia.gpu
      repository: pytorch/vision
      gpu-arch-type: cuda
      gpu-arch-version: ${{ matrix.cuda_arch_version }}
      timeout: 120
      script: |
        # Mark Build Directory Safe
        git config --global --add safe.directory /__w/vision/vision

        # Set up Environment Variables
        export PYTHON_VERSION="${{ matrix.python_version }}"
        export VERSION="${{ matrix.cuda_arch_version }}"
        export CUDATOOLKIT="pytorch-cuda=${VERSION}"

        # Set CHANNEL
        if [[ (${GITHUB_EVENT_NAME} = 'pull_request' && (${GITHUB_BASE_REF} = 'release'*)) || (${GITHUB_REF} = 'refs/heads/release'*) ]]; then
          export CHANNEL=test
        else
          export CHANNEL=nightly
        fi

        # Create Conda Env
        conda create -yp ci_env python="${PYTHON_VERSION}" numpy libpng jpeg scipy
        conda activate /work/ci_env
        
        # Install PyTorch, Torchvision, and testing libraries
        set -ex
        conda install \
          --yes \
          -c "pytorch-${CHANNEL}" \
          -c nvidia "pytorch-${CHANNEL}"::pytorch[build="*${VERSION}*"] \
          "${CUDATOOLKIT}"
        python3 setup.py develop
        python3 -m pip install pytest pytest-mock 'av<10'

        # Run Tests
        python3 -m torch.utils.collect_env
        python3 -m pytest --junitxml=test-results/junit.xml -v --durations 20