DemoModel

Author: Borda

src/litmodels/demos/__init__.py

@@ -5,10 +5,10 @@
 __all__ = []
 
 if module_available("lightning"):
-    from lightning.pytorch.demos.boring_classes import BoringModel
+    from lightning.pytorch.demos.boring_classes import BoringModel, DemoModel
 
-    __all__ += ["BoringModel"]
+    __all__ += ["BoringModel", "DemoModel"]
 elif module_available("pytorch_lightning"):
-    from pytorch_lightning.demos.boring_classes import BoringModel
+    from pytorch_lightning.demos.boring_classes import BoringModel, DemoModel
 
-    __all__ += ["BoringModel"]
+    __all__ += ["BoringModel", "DemoModel"]