colossalai

1
{
2
  "build": [
3
    {
4
      "torch_command": "pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121",
5
      "cuda_image": "hpcaitech/cuda-conda:12.1"
6
    },
7
    {
8
      "torch_command": "pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118",
9
      "cuda_image": "hpcaitech/cuda-conda:11.8"
10
    },
11
    {
12
      "torch_command": "pip install torch==2.0.0 torchvision==0.15.1 torchaudio==2.0.1",
13
      "cuda_image": "hpcaitech/cuda-conda:11.7"
14
    }
15
  ]
16
}
17