Files
ray-project--ray/python/ray/train/tensorflow/__init__.py
T
2026-07-13 13:17:40 +08:00

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.