diff --git a/colossalai/amp/apex_amp/__init__.py b/colossalai/amp/apex_amp/__init__.py
index 6689a157c..51b9b97dc 100644
--- a/colossalai/amp/apex_amp/__init__.py
+++ b/colossalai/amp/apex_amp/__init__.py
@@ -1,7 +1,8 @@
-from .apex_amp import ApexAMPOptimizer
 import torch.nn as nn
 from torch.optim import Optimizer
 
+from .apex_amp import ApexAMPOptimizer
+
 
 def convert_to_apex_amp(model: nn.Module, optimizer: Optimizer, amp_config):
     r"""A helper function to wrap training components with Apex AMP modules