Update

1 files changed, 1 insertions, 4 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 48921d4..f4d4cbb 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -18,7 +18,6 @@ import transformers
 from util.files import load_config, load_embeddings_from_dir
 from data.csv import VlpnDataModule, keyword_filter
 from training.functional import train, get_models
-from training.lr import plot_metrics
 from training.strategy.dreambooth import dreambooth_strategy
 from training.optimization import get_scheduler
 from training.util import save_args
@@ -692,7 +691,7 @@ def main():
         mid_point=args.lr_mid_point,
     )
 
-    metrics = trainer(
+    trainer(
         strategy=dreambooth_strategy,
         project="dreambooth",
         train_dataloader=datamodule.train_dataloader,
@@ -721,8 +720,6 @@ def main():
         sample_image_size=args.sample_image_size,
     )
 
-    plot_metrics(metrics, output_dir / "lr.png")
-
 
 if __name__ == "__main__":
     main()