-
Notifications
You must be signed in to change notification settings - Fork 245
/
__init__.py
34 lines (27 loc) · 981 Bytes
/
__init__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
"""
.. currentmodule:: flytekit.extras.tensorflow
.. autosummary::
:template: custom.rst
:toctree: generated/
TensorFlowRecordFileTransformer
TensorFlowRecordsDirTransformer
"""
from flytekit.loggers import logger
# TODO: abstract this out so that there's an established pattern for registering plugins
# that have soft dependencies
try:
# isolate the exception to the tensorflow import
import tensorflow
_tensorflow_installed = True
except TypeError as e:
logger.warn(f"Unsupported version of tensorflow installed. Error message from protobuf library: {e}")
_tensorflow_installed = False
except (ImportError, OSError):
_tensorflow_installed = False
if _tensorflow_installed:
from .record import TensorFlowRecordFileTransformer, TensorFlowRecordsDirTransformer
else:
logger.info(
"We won't register TensorFlowRecordFileTransformer and TensorFlowRecordsDirTransformer "
"because tensorflow is not installed."
)