diff --git a/requirements.txt b/requirements.txt index 05c11f1..d4d06d1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ # NOTE: once we add more dependencies, consider update dependabot to check for updates lightning-sdk >=0.1.35 +lightning-utilities diff --git a/src/litmodels/demos/__init__.py b/src/litmodels/demos/__init__.py new file mode 100644 index 0000000..d0d0f93 --- /dev/null +++ b/src/litmodels/demos/__init__.py @@ -0,0 +1,14 @@ +"""Define a demos model for examples and testing purposes.""" + +from lightning_utilities import module_available + +__all__ = [] + +if module_available("lightning"): + from lightning.pytorch.demos.boring_classes import BoringModel, DemoModel + + __all__ += ["BoringModel", "DemoModel"] +elif module_available("pytorch_lightning"): + from pytorch_lightning.demos.boring_classes import BoringModel, DemoModel + + __all__ += ["BoringModel", "DemoModel"]