From d5696615a84a768307e82d13e50b4aef64f69dbd Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Mon, 16 Jan 2023 10:03:05 +0100
Subject: Extended Dreambooth: Train TI tokens separately

---
 training/functional.py | 1 +
 1 file changed, 1 insertion(+)

(limited to 'training')

diff --git a/training/functional.py b/training/functional.py
index f5c111e..1b6162b 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -486,6 +486,7 @@ def train_loop(
         if accelerator.is_main_process:
             print("Interrupted")
             on_checkpoint(global_step + global_step_offset, "end")
+        raise KeyboardInterrupt
 
 
 def train(
-- 
cgit v1.2.3-70-g09d2