Set TF_NCCL_VERSION.
(environ_cp)
| 946 | |
| 947 | |
| 948 | def set_tf_nccl_version(environ_cp): |
| 949 | """Set TF_NCCL_VERSION.""" |
| 950 | if not is_linux(): |
| 951 | raise ValueError('Currently NCCL is only supported on Linux platform.') |
| 952 | |
| 953 | if 'TF_NCCL_VERSION' in environ_cp: |
| 954 | return |
| 955 | |
| 956 | ask_nccl_version = ( |
| 957 | 'Please specify the locally installed NCCL version you want to use. ' |
| 958 | '[Leave empty to use http://github.com/nvidia/nccl]: ') |
| 959 | tf_nccl_version = get_from_env_or_user_or_default(environ_cp, |
| 960 | 'TF_NCCL_VERSION', |
| 961 | ask_nccl_version, '') |
| 962 | environ_cp['TF_NCCL_VERSION'] = tf_nccl_version |
| 963 | |
| 964 | |
| 965 | def get_native_cuda_compute_capabilities(environ_cp): |
no test coverage detected