👕 Add TU trainer, CLI & datamodule to root init

ENSTA-U2IS-AI · Oct 17, 2024 · 3435604 · 3435604
1 parent 7e7933a
commit 3435604
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 0 deletions.
diff --git a/.gitignore b/.gitignore
@@ -3,6 +3,7 @@
 data/
 logs/
 lightning_logs/
+auto_tutorials_source/*.png
 docs/*/generated/
 docs/*/auto_tutorials/
 *.pth

diff --git a/torch_uncertainty/__init__.py b/torch_uncertainty/__init__.py
@@ -0,0 +1,3 @@
+# ruff: noqa: F401
+from datamodules import TUDataModule
+from utils import TULightningCLI, TUTrainer
diff --git a/torch_uncertainty/datamodules/__init__.py b/torch_uncertainty/datamodules/__init__.py
@@ -1,4 +1,5 @@
 # ruff: noqa: F401
+from .abstract import TUDataModule
 from .classification.cifar10 import CIFAR10DataModule
 from .classification.cifar100 import CIFAR100DataModule
 from .classification.imagenet import ImageNetDataModule