3 files changed, 3 insertions, 1 deletions

diff --git a/data/csv.py b/data/csv.py
index 233f5d8..619452e 100644
--- a/data/csv.py
+++ b/data/csv.py
@@ -403,6 +403,7 @@ class VlpnDataset(IterableDataset):
             if len(batch) >= batch_size:
                 yield batch
                 batch = []
+                continue
 
             bucket_mask = mask.logical_and(self.bucket_assignments == bucket)
             bucket_items = self.bucket_items[bucket_mask]
diff --git a/train_ti.py b/train_ti.py
index 171d085..f78c7d2 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -586,6 +586,7 @@ def main():
         seed=args.seed,
         with_prior_preservation=args.num_class_images != 0,
         prior_loss_weight=args.prior_loss_weight,
+        low_freq_noise=0,
         strategy=textual_inversion_strategy,
         num_train_epochs=args.num_train_epochs,
         sample_frequency=args.sample_frequency,
diff --git a/training/functional.py b/training/functional.py
index a9c7a8a..e1035ce 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -556,7 +556,7 @@ def train(
     global_step_offset: int = 0,
     with_prior_preservation: bool = False,
     prior_loss_weight: float = 1.0,
-    low_freq_noise: float = 0.05,
+    low_freq_noise: float = 0.1,
     **kwargs,
 ):
     text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, extra = strategy.prepare(