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

pull/1849/head
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

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
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,