@@ -31,7 +31,10 @@ def arch_type(arch_version: str) -> str:
3131WHEEL_CONTAINER_IMAGES = {
3232 ** {
3333 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"
3538 },
3639 ** {
3740 gpu_arch : f"pytorch/manylinux-builder:rocm{ gpu_arch } "
@@ -219,14 +222,16 @@ def generate_wheels_matrix(os: str,
219222 "container_image" : WHEEL_CONTAINER_IMAGES [arch_version ],
220223 "package_type" : package_type ,
221224 "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' | "
223227 "nvidia-cuda-cupti-cu11==11.7.101; platform_system == 'Linux' | "
224228 "nvidia-cudnn-cu11==8.5.0.96; platform_system == 'Linux' | "
225229 "nvidia-cublas-cu11==11.10.3.66; platform_system == 'Linux' | "
226230 "nvidia-cufft-cu11==10.9.0.58; platform_system == 'Linux' | "
227231 "nvidia-curand-cu11==10.2.10.91; platform_system == 'Linux' | "
228232 "nvidia-cusolver-cu11==11.4.0.1; platform_system == 'Linux' | "
229233 "nvidia-cusparse-cu11==11.7.4.91; platform_system == 'Linux' | "
234+ "nvidia-nccl-cu11==2.14.3; platform_system == 'Linux' | "
230235 "nvidia-nvtx-cu11==11.7.91; platform_system == 'Linux'" ,
231236 "build_name" :
232237 f"{ package_type } -py{ python_version } -{ gpu_arch_type } { gpu_arch_version } -with-pypi-cudnn"
0 commit comments