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

pull/1849/head
Genghan Zhang 2 years ago committed by binmakeswell
parent b0a138aa22
commit b25030cc07

@ -1,9 +1,12 @@
from typing import Optional
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.optim import Optimizer
from colossalai.context import Config from colossalai.context import Config
from .torch_amp import TorchAMPOptimizer, TorchAMPModel, TorchAMPLoss
from typing import Optional from .torch_amp import TorchAMPLoss, TorchAMPModel, TorchAMPOptimizer
def convert_to_torch_amp(model: nn.Module, def convert_to_torch_amp(model: nn.Module,

Loading…
Cancel
Save