github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-11 04:41:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: siqi <siqi@siqis-MacBook-Pro.local>
This commit is contained in:
Sze-qq 2023-03-28 15:55:09 +08:00 committed by binmakeswell
parent 8af977f223
commit 94eec1c5ad
1 changed files with 7 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
colossalai/engine/gradient_accumulation/_gradient_accumulation.py
View File

@ -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):