github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-20 02:10:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Genghan Zhang 2022-11-08 19:35:19 +08:00 committed by binmakeswell
parent b0a138aa22
commit b25030cc07
1 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
colossalai/amp/torch_amp/__init__.py
View File

@ -1,10 +1,13 @@
import torch.nn as nn
from torch.optim import Optimizer
from torch.nn.modules.loss import _Loss
from colossalai.context import Config
from .torch_amp import TorchAMPOptimizer, TorchAMPModel, TorchAMPLoss
from typing import Optional from typing import Optional
import torch.nn as nn
from torch.nn.modules.loss import _Loss
from torch.optim import Optimizer
from colossalai.context import Config
from .torch_amp import TorchAMPLoss, TorchAMPModel, TorchAMPOptimizer
def convert_to_torch_amp(model: nn.Module, def convert_to_torch_amp(model: nn.Module,
optimizer: Optimizer, optimizer: Optimizer,