From 7a06a58f49780d6938dc923c13b491f1091596bd Mon Sep 17 00:00:00 2001
From: Artiprocher <wangye87v5@hotmail.com>
Date: Fri, 11 Apr 2025 16:23:14 +0800
Subject: [PATCH] support reference image

---
 train_flux_reference.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/train_flux_reference.py b/train_flux_reference.py
index 29670be..18b47a6 100644
--- a/train_flux_reference.py
+++ b/train_flux_reference.py
@@ -212,6 +212,7 @@ if __name__ == '__main__':
             ),
         ],
         dataset_weight=(4, 2, 2, 1),
+        steps_per_epoch=args.steps_per_epoch,
     )
     train_loader = torch.utils.data.DataLoader(
         dataset,