Assign unused images in validation dataset to train dataset

1 files changed, 9 insertions, 3 deletions

diff --git a/textual_inversion.py b/textual_inversion.py
index 7a7d7fc..fa6214e 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -604,9 +604,15 @@ def main():
     )
 
     datamodule = CSVDataModule(
-        data_file=args.train_data_file, batch_size=args.train_batch_size, tokenizer=tokenizer,
-        size=args.resolution, placeholder_token=args.placeholder_token, repeats=args.repeats,
-        center_crop=args.center_crop)
+        data_file=args.train_data_file,
+        batch_size=args.train_batch_size,
+        tokenizer=tokenizer,
+        size=args.resolution,
+        placeholder_token=args.placeholder_token,
+        repeats=args.repeats,
+        center_crop=args.center_crop,
+        valid_set_size=args.sample_batch_size*args.stable_sample_batches
+    )
 
     datamodule.prepare_data()
     datamodule.setup()