Reverted modularization mostly

1 files changed, 2 insertions, 1 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index c892ebf..2145e2b 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -21,7 +21,8 @@ from slugify import slugify
 from util import load_config, load_embeddings_from_dir
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from data.csv import VlpnDataModule, VlpnDataItem
-from training.common import loss_step, train_loop, generate_class_images, get_scheduler
+from training.common import loss_step, train_loop, generate_class_images
+from training.optimization import get_scheduler
 from training.lr import LRFinder
 from training.util import CheckpointerBase, save_args
 from models.clip.embeddings import patch_managed_embeddings