[nfc]fix ColossalaiOptimizer is not defined (#4122)

09fe9dc7 · digger yu · GitHub · edd75a59 · 09fe9dc7
Unverified Commit 09fe9dc7 authored Jun 30, 2023 by digger yu Committed by GitHub Jun 30, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

colossalai/engine/_base_engine.py colossalai/engine/_base_engine.py +1 -1

No files found.
--- a/colossalai/engine/_base_engine.py
+++ b/colossalai/engine/_base_engine.py
@@ -12,7 +12,7 @@ from colossalai.engine.gradient_handler import BaseGradientHandler
 from colossalai.engine.schedule import BaseSchedule, InterleavedPipelineSchedule, NonPipelineSchedule, PipelineSchedule
 from colossalai.logging import get_dist_logger
 from colossalai.zero.legacy.gemini import BaseOpHook, register_ophooks_recursively
-
+from colossalai.nn.optimizer import ColossalaiOptimizer

 class Engine:
    """Basic engine class for training and evaluation. It runs a specific process method