mirror of https://github.com/hpcaitech/ColossalAI
[NFC] polish ./colossalai/amp/torch_amp/__init__.py code style (#1836)
parent
b0a138aa22
commit
b25030cc07
|
@ -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
|
||||
|
||||
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,
|
||||
optimizer: Optimizer,
|
||||
|
|
Loading…
Reference in New Issue