From be3e05e47cded8487aaa787c54aa74770f9dcac8 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Fri, 31 Mar 2023 09:34:55 +0200 Subject: Fix --- train_ti.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) (limited to 'train_ti.py') diff --git a/train_ti.py b/train_ti.py index 9c4ad93..b7ea5f3 100644 --- a/train_ti.py +++ b/train_ti.py @@ -761,7 +761,8 @@ def main(): num_train_epochs = args.num_train_epochs if num_train_epochs is None: - num_train_epochs = math.ceil(args.num_train_steps / len(datamodule.train_dataset)) + num_images = math.ceil(len(datamodule.train_dataset) / args.train_batch_size) * args.train_batch_size + num_train_epochs = math.ceil(args.num_train_steps / num_images) optimizer = create_optimizer( text_encoder.text_model.embeddings.temp_token_embedding.parameters(), -- cgit v1.2.3-54-g00ecf