From f34dab4270bf18fd4b830faf289d4bba254207d5 Mon Sep 17 00:00:00 2001 From: Jiarui Fang Date: Fri, 28 Oct 2022 14:48:54 +0800 Subject: [PATCH] [compatibility] ChunkMgr import error (#1772) --- colossalai/gemini/__init__.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/colossalai/gemini/__init__.py b/colossalai/gemini/__init__.py index a82640d67..9c7407eb5 100644 --- a/colossalai/gemini/__init__.py +++ b/colossalai/gemini/__init__.py @@ -1,6 +1,8 @@ -from .chunk import TensorInfo, TensorState +from .chunk import ChunkManager, TensorInfo, TensorState +from .gemini_mgr import GeminiManager from .stateful_tensor_mgr import StatefulTensorMgr from .tensor_placement_policy import TensorPlacementPolicyFactory -from .gemini_mgr import GeminiManager -__all__ = ['StatefulTensorMgr', 'TensorPlacementPolicyFactory', 'GeminiManager', 'TensorInfo', 'TensorState'] +__all__ = [ + 'StatefulTensorMgr', 'TensorPlacementPolicyFactory', 'GeminiManager', 'TensorInfo', 'TensorState', 'ChunkManager' +]