github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-23 02:20:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

[format] Run lint on colossalai.engine (#3367)

Browse Source
This commit is contained in:
Hakjin Lee
2023-04-06 00:24:43 +09:00
committed by GitHub
parent b92313903f
commit 46c009dba4
9 changed files with 32 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
colossalai/engine/gradient_accumulation/__init__.py
View File

@@ -1,10 +1,17 @@
from typing import Iterable, List
import torch.nn as nn
from typing import List
from colossalai.engine import BaseGradientHandler
from typing import Iterable
from torch.optim import Optimizer
from torch.optim.lr_scheduler import _LRScheduler
from ._gradient_accumulation import GradAccumDataloader, GradAccumOptimizer, GradAccumLrSchedulerByStep, GradAccumGradientHandler
from colossalai.engine import BaseGradientHandler
from ._gradient_accumulation import (
GradAccumDataloader,
GradAccumGradientHandler,
GradAccumLrSchedulerByStep,
GradAccumOptimizer,
)
__all__ = [
'accumulate_gradient', 'GradAccumDataloader', 'GradAccumOptimizer', 'GradAccumLrSchedulerByStep',