1 files changed, 3 insertions, 3 deletions

diff --git a/train_ti.py b/train_ti.py
index fc0d68c..344b412 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -750,6 +750,7 @@ def main():
         accelerator=accelerator,
         unet=unet,
         text_encoder=text_encoder,
+        tokenizer=tokenizer,
         vae=vae,
         noise_scheduler=noise_scheduler,
         dtype=weight_dtype,
@@ -764,8 +765,6 @@ def main():
         global_step_offset=global_step_offset,
         offset_noise_strength=args.offset_noise_strength,
         # --
-        tokenizer=tokenizer,
-        sample_scheduler=sample_scheduler,
         checkpoint_output_dir=checkpoint_output_dir,
         use_emb_decay=args.use_emb_decay,
         emb_decay_target=args.emb_decay_target,
@@ -774,6 +773,7 @@ def main():
         ema_inv_gamma=args.ema_inv_gamma,
         ema_power=args.ema_power,
         ema_max_decay=args.ema_max_decay,
+        sample_scheduler=sample_scheduler,
         sample_batch_size=args.sample_batch_size,
         sample_num_batches=args.sample_batches,
         sample_num_steps=args.sample_steps,
@@ -863,9 +863,9 @@ def main():
             optimizer=optimizer,
             lr_scheduler=lr_scheduler,
             num_train_epochs=num_train_epochs,
-            sample_frequency=sample_frequency,
             # --
             sample_output_dir=sample_output_dir,
+            sample_frequency=sample_frequency,
             placeholder_tokens=placeholder_tokens,
             placeholder_token_ids=placeholder_token_ids,
         )