Set TF_TENSORRT_VERSION.
(environ_cp)
| 929 | |
| 930 | |
| 931 | def set_tf_tensorrt_version(environ_cp): |
| 932 | """Set TF_TENSORRT_VERSION.""" |
| 933 | if not is_linux(): |
| 934 | raise ValueError('Currently TensorRT is only supported on Linux platform.') |
| 935 | |
| 936 | if not int(environ_cp.get('TF_NEED_TENSORRT', False)): |
| 937 | return |
| 938 | |
| 939 | ask_tensorrt_version = ( |
| 940 | 'Please specify the TensorRT version you want to use. ' |
| 941 | '[Leave empty to default to TensorRT %s]: ') % _DEFAULT_TENSORRT_VERSION |
| 942 | tf_tensorrt_version = get_from_env_or_user_or_default( |
| 943 | environ_cp, 'TF_TENSORRT_VERSION', ask_tensorrt_version, |
| 944 | _DEFAULT_TENSORRT_VERSION) |
| 945 | environ_cp['TF_TENSORRT_VERSION'] = tf_tensorrt_version |
| 946 | |
| 947 | |
| 948 | def set_tf_nccl_version(environ_cp): |
no test coverage detected