[format] Run lint on colossalai.engine (#3367)

colossalai/engine/schedule/__init__.py

@@ -1,5 +1,5 @@
 from ._base_schedule import BaseSchedule
-from ._pipeline_schedule import PipelineSchedule, InterleavedPipelineSchedule, get_tensor_shape
 from ._non_pipeline_schedule import NonPipelineSchedule
+from ._pipeline_schedule import InterleavedPipelineSchedule, PipelineSchedule, get_tensor_shape
 
 __all__ = ['BaseSchedule', 'NonPipelineSchedule', 'PipelineSchedule', 'InterleavedPipelineSchedule', 'get_tensor_shape']

colossalai/engine/schedule/_base_schedule.py

@@ -2,10 +2,10 @@
 # -*- encoding: utf-8 -*-
 
 from abc import ABC, abstractmethod
+from typing import Callable, Iterable
 
 import torch
 
-from typing import Iterable, Callable
 from colossalai.logging import get_dist_logger
 from colossalai.utils import get_current_device
 

colossalai/engine/schedule/_non_pipeline_schedule.py

@@ -1,13 +1,14 @@
 #!/usr/bin/env python
 # -*- encoding: utf-8 -*-
 
-from typing import Iterable
+import inspect
+from typing import Callable, Iterable
 
 import torch
-import inspect
-from ._base_schedule import BaseSchedule
+
 from colossalai.utils import conditional_context
-from typing import Callable
+
+from ._base_schedule import BaseSchedule
 
 
 class NonPipelineSchedule(BaseSchedule):