@@ -31,7 +31,10 @@ def arch_type(arch_version: str) -> str:
31
31
WHEEL_CONTAINER_IMAGES = {
32
32
** {
33
33
gpu_arch : f"pytorch/manylinux-builder:cuda{ gpu_arch } "
34
- for gpu_arch in CUDA_ARCHES
34
+ for gpu_arch in CUDA_ARCHES if gpu_arch == "11.6"
35
+ },
36
+ ** {
37
+ "11.7" : "tousif111/manylinux-builder:cuda11.7"
35
38
},
36
39
** {
37
40
gpu_arch : f"pytorch/manylinux-builder:rocm{ gpu_arch } "
@@ -219,14 +222,16 @@ def generate_wheels_matrix(os: str,
219
222
"container_image" : WHEEL_CONTAINER_IMAGES [arch_version ],
220
223
"package_type" : package_type ,
221
224
"pytorch_extra_install_requirements" :
222
- "nvidia-cuda-runtime-cu11; platform_system == 'Linux' | "
225
+ "nvidia-cuda-nvrtc-cu11==11.7.99; platform_system == 'Linux' | "
226
+ "nvidia-cuda-runtime-cu11==11.7.99; platform_system == 'Linux' | "
223
227
"nvidia-cuda-cupti-cu11==11.7.101; platform_system == 'Linux' | "
224
228
"nvidia-cudnn-cu11==8.5.0.96; platform_system == 'Linux' | "
225
229
"nvidia-cublas-cu11==11.10.3.66; platform_system == 'Linux' | "
226
230
"nvidia-cufft-cu11==10.9.0.58; platform_system == 'Linux' | "
227
231
"nvidia-curand-cu11==10.2.10.91; platform_system == 'Linux' | "
228
232
"nvidia-cusolver-cu11==11.4.0.1; platform_system == 'Linux' | "
229
233
"nvidia-cusparse-cu11==11.7.4.91; platform_system == 'Linux' | "
234
+ "nvidia-nccl-cu11==2.14.3; platform_system == 'Linux' | "
230
235
"nvidia-nvtx-cu11==11.7.91; platform_system == 'Linux'" ,
231
236
"build_name" :
232
237
f"{ package_type } -py{ python_version } -{ gpu_arch_type } { gpu_arch_version } -with-pypi-cudnn"
0 commit comments