@@ -235,25 +235,27 @@ def tests(session: Session) -> None:
235235 install_additional_warning = """
236236 Installation failed, this likely means that the required hardware or drivers are missing.
237237 Run without `-- additional-tests` to avoid this."""
238+ # install cuda-python
238239 try :
239240 session .install ("cuda-python" )
240241 except Exception as error :
241242 session .log (error )
242243 session .warn (install_additional_warning )
244+ # install cupy
245+ try :
243246 try :
244- try :
245- # based on the CUDA version, try installing the exact prebuilt cupy version
246- cuda_cupy_version = f"cupy-cuda{ '' .join (cuda_version .split ('.' ))} "
247- session .install (cuda_cupy_version )
248- except Exception :
249- # if the exact prebuilt is not available, try the more general prebuilt
250- cuda_cupy_version_x = f"cupy-cuda{ cuda_version .split ('.' )[0 ]} x"
251- session .warn (f"CuPy exact prebuilt not available for { cuda_version } , trying { cuda_cupy_version_x } " )
252- session .install (cuda_cupy_version_x )
247+ # based on the CUDA version, try installing the exact prebuilt cupy version
248+ cuda_cupy_version = f"cupy-cuda{ '' .join (cuda_version .split ('.' ))} "
249+ session .install (cuda_cupy_version )
253250 except Exception :
254- # if no compatible prebuilt wheel is found, try building CuPy ourselves
255- session .warn (f"No prebuilt CuPy found for CUDA { cuda_version } , building from source..." )
256- session .install ("cupy" )
251+ # if the exact prebuilt is not available, try the more general prebuilt
252+ cuda_cupy_version_x = f"cupy-cuda{ cuda_version .split ('.' )[0 ]} x"
253+ session .warn (f"CuPy exact prebuilt not available for { cuda_version } , trying { cuda_cupy_version_x } " )
254+ session .install (cuda_cupy_version_x )
255+ except Exception :
256+ # if no compatible prebuilt wheel is found, try building CuPy ourselves
257+ session .warn (f"No prebuilt CuPy found for CUDA { cuda_version } , building from source..." )
258+ session .install ("cupy" )
257259
258260 # for the last Python version session if all optional dependencies are enabled:
259261 if session .python == python_versions_to_test [- 1 ] and full_install :
0 commit comments