From b59611506d2eac1b280ada07db0bf9de15b41a26 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Tue, 28 Mar 2023 16:28:44 +0200 Subject: Fix --- train_dreambooth.py | 2 +- train_lora.py | 2 +- train_ti.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/train_dreambooth.py b/train_dreambooth.py index acb8287..e3c8525 100644 --- a/train_dreambooth.py +++ b/train_dreambooth.py @@ -598,7 +598,7 @@ def main(): num_train_epochs = args.num_train_epochs if num_train_epochs is None: - num_train_epochs = math.ceil(len(datamodule.train_dataset) / args.num_train_steps) + num_train_epochs = math.ceil(args.num_train_steps / len(datamodule.train_dataset)) params_to_optimize = (unet.parameters(), ) if args.train_text_encoder_epochs != 0: diff --git a/train_lora.py b/train_lora.py index a9c6e52..6f8644b 100644 --- a/train_lora.py +++ b/train_lora.py @@ -630,7 +630,7 @@ def main(): num_train_epochs = args.num_train_epochs if num_train_epochs is None: - num_train_epochs = math.ceil(len(datamodule.train_dataset) / args.num_train_steps) + num_train_epochs = math.ceil(args.num_train_steps / len(datamodule.train_dataset)) optimizer = create_optimizer( itertools.chain( diff --git a/train_ti.py b/train_ti.py index 7bcc72f..9c4ad93 100644 --- a/train_ti.py +++ b/train_ti.py @@ -761,7 +761,7 @@ def main(): num_train_epochs = args.num_train_epochs if num_train_epochs is None: - num_train_epochs = math.ceil(len(datamodule.train_dataset) / args.num_train_steps) + num_train_epochs = math.ceil(args.num_train_steps / len(datamodule.train_dataset)) optimizer = create_optimizer( text_encoder.text_model.embeddings.temp_token_embedding.parameters(), -- cgit v1.2.3-70-g09d2