diff --git a/library/train_util.py b/library/train_util.py
index 1b4a29c3f..34ea21589 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -19,6 +19,7 @@
     Sequence,
     Tuple,
     Union,
+    Callable,
 )
 from accelerate import Accelerator, InitProcessGroupKwargs, DistributedDataParallelKwargs, PartialState
 import glob