Assign unused images in validation dataset to train dataset

1 files changed, 2 insertions, 1 deletions

diff --git a/dreambooth.py b/dreambooth.py
index 744d1bc..88cd0da 100644
--- a/dreambooth.py
+++ b/dreambooth.py
@@ -638,6 +638,7 @@ def main():
         identifier=args.identifier,
         repeats=args.repeats,
         center_crop=args.center_crop,
+        valid_set_size=args.sample_batch_size*args.stable_sample_batches,
         collate_fn=collate_fn)
 
     datamodule.prepare_data()
@@ -658,7 +659,7 @@ def main():
         sample_batch_size=args.sample_batch_size,
         random_sample_batches=args.random_sample_batches,
         stable_sample_batches=args.stable_sample_batches,
-        seed=args.seed
+        seed=args.seed or torch.random.seed()
     )
 
     # Scheduler and math around the number of training steps.