31 lines
865 B
Python
31 lines
865 B
Python
# isort: off
|
|
try:
|
|
import tensorflow as tf # noqa: F401
|
|
except ModuleNotFoundError:
|
|
raise ModuleNotFoundError(
|
|
"TensorFlow isn't installed. To install TensorFlow, run 'pip install "
|
|
"tensorflow'."
|
|
)
|
|
# isort: on
|
|
|
|
from ray.train.tensorflow.config import TensorflowConfig
|
|
from ray.train.tensorflow.tensorflow_checkpoint import TensorflowCheckpoint
|
|
from ray.train.tensorflow.tensorflow_trainer import TensorflowTrainer
|
|
from ray.train.tensorflow.train_loop_utils import prepare_dataset_shard
|
|
from ray.train.v2._internal.constants import is_v2_enabled
|
|
|
|
if is_v2_enabled():
|
|
from ray.train.v2.tensorflow.tensorflow_trainer import ( # noqa: F811
|
|
TensorflowTrainer,
|
|
)
|
|
|
|
__all__ = [
|
|
"TensorflowCheckpoint",
|
|
"TensorflowConfig",
|
|
"prepare_dataset_shard",
|
|
"TensorflowTrainer",
|
|
]
|
|
|
|
|
|
# DO NOT ADD ANYTHING AFTER THIS LINE.
|