From cb5a587e9aa545a41980ee68e88bf5edf59c44cb Mon Sep 17 00:00:00 2001 From: Jiarui Fang Date: Wed, 2 Nov 2022 12:10:52 +0800 Subject: [PATCH] [hotfix] polish chunk import (#1787) --- colossalai/gemini/__init__.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/colossalai/gemini/__init__.py b/colossalai/gemini/__init__.py index 9c7407eb5..7a5a44ebb 100644 --- a/colossalai/gemini/__init__.py +++ b/colossalai/gemini/__init__.py @@ -1,8 +1,9 @@ -from .chunk import ChunkManager, TensorInfo, TensorState +from .chunk import ChunkManager, TensorInfo, TensorState, search_chunk_configuration from .gemini_mgr import GeminiManager from .stateful_tensor_mgr import StatefulTensorMgr from .tensor_placement_policy import TensorPlacementPolicyFactory __all__ = [ - 'StatefulTensorMgr', 'TensorPlacementPolicyFactory', 'GeminiManager', 'TensorInfo', 'TensorState', 'ChunkManager' + 'StatefulTensorMgr', 'TensorPlacementPolicyFactory', 'GeminiManager', 'TensorInfo', 'TensorState', 'ChunkManager', + 'search_chunk_configuration' ]