From 44b55445149f61d140daedbce6b0469128468dd4 Mon Sep 17 00:00:00 2001
From: "baodong.lh" <baodong.lh@alibaba-inc.com>
Date: Fri, 27 Dec 2024 23:14:43 +0800
Subject: [PATCH] fix generator call

---
 chatlearn/utils/dist_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/chatlearn/utils/dist_utils.py b/chatlearn/utils/dist_utils.py
index c3e65dc4..e5976a3d 100644
--- a/chatlearn/utils/dist_utils.py
+++ b/chatlearn/utils/dist_utils.py
@@ -33,7 +33,7 @@ def bucket_tensor_generator(tensor_generator, bucket_size_mb):
     """
     size_limit = bucket_size_mb * 1024 * 1024
     buf_dict = defaultdict(lambda: [[], 0])
-    for tensor in tensor_generator:
+    for tensor in tensor_generator():
         if tensor.is_sparse:
             yield tensor, False
             continue