IAM datasets implemented.

1 files changed, 13 insertions, 2 deletions

diff --git a/src/training/trainer/callbacks/__init__.py b/src/training/trainer/callbacks/__init__.py
index 5942276..c81e4bf 100644
--- a/src/training/trainer/callbacks/__init__.py
+++ b/src/training/trainer/callbacks/__init__.py
@@ -1,7 +1,16 @@
 """The callback modules used in the training script."""
-from .base import Callback, CallbackList, Checkpoint
+from .base import Callback, CallbackList
+from .checkpoint import Checkpoint
 from .early_stopping import EarlyStopping
-from .lr_schedulers import CyclicLR, MultiStepLR, OneCycleLR, ReduceLROnPlateau, StepLR
+from .lr_schedulers import (
+    CosineAnnealingLR,
+    CyclicLR,
+    MultiStepLR,
+    OneCycleLR,
+    ReduceLROnPlateau,
+    StepLR,
+    SWA,
+)
 from .progress_bar import ProgressBar
 from .wandb_callbacks import WandbCallback, WandbImageLogger
 
@@ -9,6 +18,7 @@ __all__ = [
     "Callback",
     "CallbackList",
     "Checkpoint",
+    "CosineAnnealingLR",
     "EarlyStopping",
     "WandbCallback",
     "WandbImageLogger",
@@ -18,4 +28,5 @@ __all__ = [
     "ProgressBar",
     "ReduceLROnPlateau",
     "StepLR",
+    "SWA",
 ]