1 files changed, 4 insertions, 2 deletions

diff --git a/train_lora.py b/train_lora.py
index 0d8b8cb..1d1485d 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -873,7 +873,6 @@ def main():
         seed=args.seed,
         guidance_scale=args.guidance_scale,
         prior_loss_weight=args.prior_loss_weight if args.num_class_images != 0 else 0,
-        offset_noise_strength=args.offset_noise_strength,
         sample_scheduler=sample_scheduler,
         sample_batch_size=args.sample_batch_size,
         sample_num_batches=args.sample_batches,
@@ -984,13 +983,14 @@ def main():
             lr_scheduler=pti_lr_scheduler,
             num_train_epochs=num_train_epochs,
             gradient_accumulation_steps=args.gradient_accumulation_steps,
-            cycle=1,
+            cycle=0,
             pti_mode=True,
             # --
             group_labels=["emb"],
             sample_output_dir=pti_sample_output_dir,
             checkpoint_output_dir=pti_checkpoint_output_dir,
             sample_frequency=pti_sample_frequency,
+            offset_noise_strength=0,
             no_val=True,
         )
 
@@ -1132,11 +1132,13 @@ def main():
             gradient_accumulation_steps=args.gradient_accumulation_steps,
             global_step_offset=training_iter * num_train_steps,
             cycle=training_iter,
+            train_text_encoder_cycles=args.train_text_encoder_cycles,
             # --
             group_labels=group_labels,
             sample_output_dir=lora_sample_output_dir,
             checkpoint_output_dir=lora_checkpoint_output_dir,
             sample_frequency=lora_sample_frequency,
+            offset_noise_strength=args.offset_noise_strength,
             no_val=args.valid_set_size == 0,
         )