From 6754f1b77fd11041a0c9db12dc4117d69a09cfac Mon Sep 17 00:00:00 2001 From: Ziyue Jiang Date: Mon, 6 Jun 2022 12:11:48 +0800 Subject: [PATCH] fix module utils bug (#1066) --- colossalai/nn/module_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/colossalai/nn/module_utils.py b/colossalai/nn/module_utils.py index 64d3f8075..158bd1062 100644 --- a/colossalai/nn/module_utils.py +++ b/colossalai/nn/module_utils.py @@ -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