__init__.py 767 B

123456789101112131415161718192021222324252627282930313233
  1. # isort: off
  2. try:
  3. import lightning # noqa: F401
  4. except ModuleNotFoundError:
  5. try:
  6. import pytorch_lightning # noqa: F401
  7. except ModuleNotFoundError:
  8. raise ModuleNotFoundError(
  9. "PyTorch Lightning isn't installed. To install PyTorch Lightning, "
  10. "please run 'pip install lightning'"
  11. )
  12. # isort: on
  13. from ray.train.lightning._lightning_utils import (
  14. RayDDPStrategy,
  15. RayDeepSpeedStrategy,
  16. RayFSDPStrategy,
  17. RayLightningEnvironment,
  18. RayTrainReportCallback,
  19. prepare_trainer,
  20. )
  21. __all__ = [
  22. "prepare_trainer",
  23. "RayDDPStrategy",
  24. "RayFSDPStrategy",
  25. "RayDeepSpeedStrategy",
  26. "RayLightningEnvironment",
  27. "RayTrainReportCallback",
  28. ]
  29. # DO NOT ADD ANYTHING AFTER THIS LINE.