Update train_util.py

gesen2egee · Mar 19, 2024 · dd4f1ed · dd4f1ed
1 parent 5f1b69e
commit dd4f1ed
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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 (