From 94329fc139f2f385a851cbfb768ddeb253631a3e Mon Sep 17 00:00:00 2001 From: LuGY <74758262+Gy-Lu@users.noreply.github.com> Date: Wed, 9 Nov 2022 14:49:42 +0800 Subject: [PATCH] [NFC] polish colossalai/amp/apex_amp/__init__.py code style (#1853) --- colossalai/amp/apex_amp/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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