Commit 94bfd351 authored by xyupeng's avatar xyupeng Committed by Frank Lee
Browse files

[NFC] polish colossalai/builder/builder.py code style (#1265)

parent db13f963
...@@ -6,7 +6,6 @@ import inspect ...@@ -6,7 +6,6 @@ import inspect
from colossalai.registry import * from colossalai.registry import *
def build_from_config(module, config: dict): def build_from_config(module, config: dict):
"""Returns an object of :class:`module` constructed from `config`. """Returns an object of :class:`module` constructed from `config`.
...@@ -46,23 +45,20 @@ def build_from_registry(config, registry: Registry): ...@@ -46,23 +45,20 @@ def build_from_registry(config, registry: Registry):
Raises: Raises:
Exception: Raises an Exception if an error occurred when building from registry. Exception: Raises an Exception if an error occurred when building from registry.
""" """
config_ = config.copy() # keep the original config untouched config_ = config.copy() # keep the original config untouched
assert isinstance( assert isinstance(registry, Registry), f'Expected type Registry but got {type(registry)}'
registry, Registry), f'Expected type Registry but got {type(registry)}'
mod_type = config_.pop('type') mod_type = config_.pop('type')
assert registry.has( assert registry.has(mod_type), f'{mod_type} is not found in registry {registry.name}'
mod_type), f'{mod_type} is not found in registry {registry.name}'
try: try:
obj = registry.get_module(mod_type)(**config_) obj = registry.get_module(mod_type)(**config_)
except Exception as e: except Exception as e:
print( print(f'An error occurred when building {mod_type} from registry {registry.name}', flush=True)
f'An error occurred when building {mod_type} from registry {registry.name}',
flush=True)
raise e raise e
return obj return obj
def build_gradient_handler(config, model, optimizer): def build_gradient_handler(config, model, optimizer):
"""Returns a gradient handler object of :class:`BaseGradientHandler` constructed from `config`, """Returns a gradient handler object of :class:`BaseGradientHandler` constructed from `config`,
`model` and `optimizer`. `model` and `optimizer`.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment