From 0540c33acac223b672da05e40edcfb3b6a35c0da Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sun, 21 Apr 2024 17:45:29 +0900 Subject: [PATCH] pop weights if available #1247 --- train_network.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/train_network.py b/train_network.py index 3a5255160..aad5a7194 100644 --- 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):