github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-22 05:29:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix module utils bug (#1066)

Browse Source
This commit is contained in:
Ziyue Jiang 2022-06-06 12:11:48 +08:00 committed by GitHub
parent a00644079e
commit 6754f1b77f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/nn/module_utils.py
View File

@ -14,7 +14,7 @@ def register_colo_module(module_type: type, colo_module: ColoModule):
def is_colo_module(module: torch.nn.Module):
global _COLOSSAL_MODULES
for module_type in _COLOSSAL_MODULES.keys():
if isinstance(type(module), module_type):
if isinstance(module, module_type):
return True
return False
@ -23,7 +23,7 @@ def get_colo_module(module: torch.nn.Module):
global _COLOSSAL_MODULES
if is_colo_module(module):
for module_type, colo_module in _COLOSSAL_MODULES.items():
if isinstance(type(module), module_type):
if isinstance(module, module_type):
return colo_module
else:
return None