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

[hotfix] fix memory leak in zero (#781)

Browse Source
This commit is contained in:
HELSON
2022-04-18 13:57:03 +08:00
committed by GitHub
parent 4b01da24cd
commit 4c4388c46e
6 changed files with 32 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
colossalai/amp/naive_amp/grad_scaler/dynamic_grad_scaler.py
View File

@@ -3,6 +3,7 @@
import torch
from .base_grad_scaler import BaseGradScaler
from typing import Optional
__all__ = ['DynamicGradScaler']
@@ -10,12 +11,12 @@ __all__ = ['DynamicGradScaler']
class DynamicGradScaler(BaseGradScaler):
def __init__(self,
initial_scale: int = 2**16,
growth_factor: int = 2,
initial_scale: float = 2**16,
growth_factor: float = 2,
backoff_factor: float = 0.5,
growth_interval: int = 1000,
min_scale: int = None,
max_scale: int = None,
min_scale: Optional[float] = None,
max_scale: Optional[float] = None,
hysteresis: int = 2,
verbose: bool = False):
super().__init__(initial_scale, verbose)