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

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