diff --git a/colossalai/engine/_base_engine.py b/colossalai/engine/_base_engine.py index 4f3654cd9..696d82e41 100644 --- a/colossalai/engine/_base_engine.py +++ b/colossalai/engine/_base_engine.py @@ -15,6 +15,7 @@ from typing import Optional, Type from colossalai.engine.gradient_handler import BaseGradientHandler from colossalai.logging import get_dist_logger + class Engine: """Basic engine class for training and evaluation. It runs a specific process method :meth:`step` which is based on the given :attr:`schedule` over each batch of a dataset.