From e1b504bca41a9793ed7e88ef14f2e2cbd85724f2 Mon Sep 17 00:00:00 2001
From: aktersnurra <gustaf.rydholm@gmail.com>
Date: Tue, 8 Sep 2020 23:14:23 +0200
Subject: IAM datasets implemented.

---
 src/training/trainer/callbacks/__init__.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

(limited to 'src/training/trainer/callbacks/__init__.py')

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",
 ]
-- 
cgit v1.2.3-70-g09d2