.cuda_ext.json

{
  "build": [
    {
      "torch_command": "pip install torch==1.12.1+cu102 torchvision==0.13.1+cu102 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu102",
      "cuda_image": "hpcaitech/cuda-conda:10.2"
    },
    {
      "torch_command": "pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113",
      "cuda_image": "hpcaitech/cuda-conda:11.3"
    },
    {
      "torch_command": "pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116",
      "cuda_image": "hpcaitech/cuda-conda:11.6"
    }
  ]
}