From 8d191215177e741f5eb70fde6c106d2f99a7bc80 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sun, 1 Sep 2024 16:19:17 +0300 Subject: [PATCH] possible resume fix --- train_network.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/train_network.py b/train_network.py index 754b19b..5f7f3ef 100644 --- a/train_network.py +++ b/train_network.py @@ -632,15 +632,15 @@ class NetworkTrainer: def save_model_hook(models, weights, output_dir): # pop weights of other models than network to save only network weights # only main process or deepspeed https://github.com/huggingface/diffusers/issues/2606 - if args.deepspeed: - remove_indices = [] - for i, model in enumerate(models): - if not isinstance(model, type(accelerator.unwrap_model(network))): - remove_indices.append(i) - for i in reversed(remove_indices): - if len(weights) > i: - weights.pop(i) - # print(f"save model hook: {len(weights)} weights will be saved") + #if args.deepspeed: + remove_indices = [] + for i, model in enumerate(models): + if not isinstance(model, type(accelerator.unwrap_model(network))): + remove_indices.append(i) + for i in reversed(remove_indices): + if len(weights) > i: + weights.pop(i) + # print(f"save model hook: {len(weights)} weights will be saved") # save current ecpoch and step train_state_file = os.path.join(output_dir, "train_state.json")