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

[hotfix] add deconstructor for stateful tensor (#848)

Browse Source 
* add deconstructor for stateful tensor

* fix colo init context
This commit is contained in:
ver217
2022-04-24 15:03:04 +08:00
committed by GitHub
parent 0f7ed8c192
commit 0dea140760
3 changed files with 20 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/utils/model/colo_init_context.py
View File

@@ -12,7 +12,7 @@ class ColoInitContext(InsertPostInitMethodToModuleSubClasses):
super().__init__()
self._lazy_memory_allocate = lazy_memory_allocate
def _post_init_method(self, module: torch.nn.Module):
def _post_init_method(self, module: torch.nn.Module, *args, **kwargs):
"""
The function to call at the end of the constructor of each module.
FIXME(fjr) The module may be passed to this function multiple times?