github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-26 12:14:02 +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

3
colossalai/amp/naive_amp/grad_scaler/base_grad_scaler.py
View File

@@ -12,7 +12,7 @@ __all__ = ['BaseGradScaler']
class BaseGradScaler(ABC):
def __init__(self, initial_scale: int, verbose: bool):
def __init__(self, initial_scale: float, verbose: bool):
assert initial_scale > 0
self._scale = torch.cuda.FloatTensor([initial_scale])
self._verbose = verbose
@@ -31,6 +31,7 @@ class BaseGradScaler(ABC):
def state_dict(self) -> Dict:
state_dict = dict()
state_dict['scale'] = self.scale
return state_dict
def load_state_dict(self, state_dict: Dict) -> None:
self._scale = state_dict['scale']