pop weights if available kohya-ss#1247

A4P7J1N7M05OT · May 12, 2024 · a3f77bf · a3f77bf
1 parent 2be64b3
commit a3f77bf
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/train_network.py b/train_network.py
@@ -481,7 +481,8 @@ def save_model_hook(models, weights, output_dir):
                     if not isinstance(model, type(accelerator.unwrap_model(network))):
                         remove_indices.append(i)
                 for i in reversed(remove_indices):
-                    weights.pop(i)
+                    if len(weights) > i:
+                        weights.pop(i)
                 # print(f"save model hook: {len(weights)} weights will be saved")
 
         def load_model_hook(models, input_dir):