[NFC] polish colossalai/engine/gradient_accumulation/_gradient_accumulation.py code style (#3277)

Co-authored-by: siqi <siqi@siqis-MacBook-Pro.local>
pull/3313/head
Sze-qq 2 years ago committed by binmakeswell
parent 8af977f223
commit 94eec1c5ad

@ -1,17 +1,18 @@
#!/usr/bin/env python #!/usr/bin/env python
# -*- encoding: utf-8 -*- # -*- encoding: utf-8 -*-
from typing import Union from typing import Any, Iterable, Tuple, Union
import torch.nn as nn import torch.nn as nn
from torch import Tensor from torch import Tensor
from typing import Iterable, Any, Tuple
from colossalai.nn.optimizer import ColossalaiOptimizer
from torch.nn.parallel.distributed import DistributedDataParallel from torch.nn.parallel.distributed import DistributedDataParallel
from torch.optim import Optimizer from torch.optim import Optimizer
from torch.optim.lr_scheduler import _LRScheduler from torch.optim.lr_scheduler import _LRScheduler
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
from colossalai.utils import conditional_context
from colossalai.engine import BaseGradientHandler from colossalai.engine import BaseGradientHandler
from colossalai.nn.optimizer import ColossalaiOptimizer
from colossalai.utils import conditional_context
class GradAccumOptimizer(ColossalaiOptimizer): class GradAccumOptimizer(ColossalaiOptimizer):

Loading…
Cancel
Save