refine script

2026-03-20 07:18:14 +00:00 · 2024-06-05 13:34:55 +08:00
parent 78f53a0754
commit 0d03cb31ad
1 changed files with 1 additions and 1 deletions
--- a/examples/hunyuan_dit/train_hunyuan_dit_lora.py
+++ b/examples/hunyuan_dit/train_hunyuan_dit_lora.py
@@ -256,7 +256,7 @@ if __name__ == '__main__':
    # dataset and data loader
    dataset = TextImageDataset(
        args.dataset_path,
-        steps_per_epoch=args.steps_per_epoch,
+        steps_per_epoch=args.steps_per_epoch * args.batch_size,
        height=args.height,
        width=args.width,
        center_crop=args.center_crop,