()
| 23 | |
| 24 | |
| 25 | def get_requirements(): |
| 26 | def get_version_range(): |
| 27 | def get_vers(var): |
| 28 | vers = os.environ.get(var).replace("cuda-", "") |
| 29 | major, minor = map(int, vers.split(".")) |
| 30 | return major, minor |
| 31 | |
| 32 | cuda_major, _ = get_vers("CUDA") |
| 33 | return "-cu{cuda_major}".format(cuda_major=cuda_major) |
| 34 | |
| 35 | reqs = ["nvidia-cuda-runtime" + get_version_range()] |
| 36 | if "##TENSORRT_MODULE##" == "tensorrt": |
| 37 | reqs += [ |
| 38 | "nvidia-cudnn" + get_version_range(), |
| 39 | "nvidia-cublas" + get_version_range(), |
| 40 | ] |
| 41 | return reqs |
| 42 | |
| 43 | |
| 44 | setup( |
no test coverage detected