pyg-lib/.github/workflows/cuda/Linux-env.sh at b5aa94cbbe4d0e88e5faf081edd52543fa3b49d1 · pyg-team/pyg-lib · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
#!/bin/bash
set -ex

case ${1} in
  cu128)
    export FORCE_CUDA=1
    export PATH=/usr/local/cuda-12.8/bin:${PATH}
    export TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;9.0;10.0;12.0+PTX"
    ;;
  cu126)
    export FORCE_CUDA=1
    export PATH=/usr/local/cuda-12.6/bin:${PATH}
    export TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6;9.0+PTX"
    ;;
  cu124)
    export FORCE_CUDA=1
    export PATH=/usr/local/cuda-12.4/bin:${PATH}
    export TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6;9.0+PTX"
    ;;
  cu121)
    export FORCE_CUDA=1
    export PATH=/usr/local/cuda-12.1/bin:${PATH}
    export TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6;9.0+PTX"
    ;;
  cu118)
    export FORCE_CUDA=1
    export PATH=/usr/local/cuda-11.8/bin:${PATH}
    export TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6;9.0+PTX"
    ;;
  cu117)
    export FORCE_CUDA=1
    export PATH=/usr/local/cuda-11.7/bin:${PATH}
    export TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6+PTX"
    ;;
  cu116)
    export FORCE_CUDA=1
    export PATH=/usr/local/cuda-11.6/bin:${PATH}
    export TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6+PTX"
    ;;
  cu115)
    export FORCE_CUDA=1
    export PATH=/usr/local/cuda-11.5/bin:${PATH}
    export TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6+PTX"
    ;;
  cu113)
    export FORCE_CUDA=1
    export PATH=/usr/local/cuda-11.3/bin:${PATH}
    export TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5;8.0;8.6+PTX"
    ;;
  cu102)
    export FORCE_CUDA=1
    export PATH=/usr/local/cuda-10.2/bin:${PATH}
    export TORCH_CUDA_ARCH_LIST="6.0;7.0;7.5+PTX"
    ;;
  *)
    ;;
esac