From 659cf3189b178e303db4f6bc79640b4a9e719da4 Mon Sep 17 00:00:00 2001
From: gesen2egee <gesen2gee@gmail.com>
Date: Sat, 5 Oct 2024 16:11:47 +0800
Subject: [PATCH] Update train_util.py

---
 library/train_util.py | 1 +
 1 file changed, 1 insertion(+)

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