From dd4f1ed80aba344e92f8e1ef1ae76464eb24698a Mon Sep 17 00:00:00 2001
From: gesen2egee <79357052+gesen2egee@users.noreply.github.com>
Date: Tue, 19 Mar 2024 15:52:27 +0800
Subject: [PATCH] Update train_util.py

---
 library/train_util.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/library/train_util.py b/library/train_util.py
index 70cd6e194..ee09aaadc 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -1533,6 +1533,7 @@ def __init__(
         validation_split: float,
         validation_seed: Optional[int],        
         debug_dataset: bool,
+        sample_weight: bool,
     ) -> None:
         super().__init__(tokenizer, max_token_length, resolution, network_multiplier, debug_dataset)
 
@@ -1545,7 +1546,7 @@ def __init__(
         self.size = min(self.width, self.height)  # 短いほう
         self.prior_loss_weight = prior_loss_weight
         self.latents_cache = None
-
+        self.sample_weight = sample_weight
         self.enable_bucket = enable_bucket
         if self.enable_bucket:
             assert (