github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-27 07:47:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish colossalai/amp/__init__.py code style (#3272)

Browse Source
This commit is contained in:
lucasliunju 2023-03-28 10:44:04 +08:00 committed by binmakeswell
parent 6b3bb2c249
commit 4b95464994
1 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
colossalai/amp/__init__.py
View File

@ -1,14 +1,16 @@
#!/usr/bin/env python #!/usr/bin/env python
# -*- encoding: utf-8 -*- # -*- encoding: utf-8 -*-
from .amp_type import AMP_TYPE
from colossalai.context import Config
import torch.nn as nn import torch.nn as nn
from torch.optim import Optimizer
from torch.nn.modules.loss import _Loss from torch.nn.modules.loss import _Loss
from .torch_amp import convert_to_torch_amp from torch.optim import Optimizer
from colossalai.context import Config
from .amp_type import AMP_TYPE
from .apex_amp import convert_to_apex_amp from .apex_amp import convert_to_apex_amp
from .naive_amp import convert_to_naive_amp from .naive_amp import convert_to_naive_amp
from .torch_amp import convert_to_torch_amp
__all__ = ['convert_to_amp', 'convert_to_naive_amp', 'convert_to_apex_amp', 'convert_to_torch_amp', 'AMP_TYPE'] __all__ = ['convert_to_amp', 'convert_to_naive_amp', 'convert_to_apex_amp', 'convert_to_torch_amp', 'AMP_TYPE']