ensure dataloader workers exit cleanly on ctrl-c

2023-01-30 14:19:18 +01:00 · 2023-01-30 14:19:18 +01:00 · c270dbf6a8
parent 4b5654452c
commit c270dbf6a8
1 changed files with 4 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -692,6 +692,9 @@ def main(args):
                time.sleep(2) # give opportunity to ctrl-C again to cancel save
                __save_model(interrupted_checkpoint_path, unet, text_encoder, tokenizer, noise_scheduler, vae, args.save_ckpt_dir, args.save_full_precision)
            exit(_SIGTERM_EXIT_CODE)
+        else:
+            # non-main threads (i.e. dataloader workers) should exit cleanly
+            exit(0)

    signal.signal(signal.SIGINT, sigterm_handler)