fix imports

Signed-off-by: Avik Basu <[email protected]>
numaproj · Aug 7, 2023 · e35559a · e35559a
1 parent 9314c8f
commit e35559a
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/numalogic/models/vae/__init__.py b/numalogic/models/vae/__init__.py
@@ -0,0 +1,3 @@
+from numalogic.models.vae.trainer import VAETrainer
+
+__all__ = ["VAETrainer"]
diff --git a/numalogic/models/vae/variants/__init__.py b/numalogic/models/vae/variants/__init__.py
@@ -0,0 +1,3 @@
+from numalogic.models.vae.variants.conv import Conv1dVAE
+
+__all__ = ["Conv1dVAE"]
diff --git a/tests/models/vae/test_conv.py b/tests/models/vae/test_conv.py
@@ -9,8 +9,8 @@
 from torch.utils.data import DataLoader
 
 from numalogic._constants import TESTS_DIR
-from numalogic.models.vae.trainer import VAETrainer
-from numalogic.models.vae.variants.conv import Conv1dVAE
+from numalogic.models.vae import VAETrainer
+from numalogic.models.vae.variants import Conv1dVAE
 from numalogic.tools.data import TimeseriesDataModule, StreamingDataset
 from numalogic.tools.exceptions import ModelInitializationError