mirror of https://github.com/hpcaitech/ColossalAI
27 lines
544 B
Python
27 lines
544 B
Python
import functools
|
|
import warnings
|
|
|
|
__all__ = ['exception_handler']
|
|
|
|
|
|
def exception_handler(func):
|
|
"""
|
|
A function wrapper to handle the AssertionError in the function.
|
|
|
|
Usage:
|
|
# mute the assertion error in the function
|
|
@exception_handler
|
|
def do_something():
|
|
...
|
|
"""
|
|
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
try:
|
|
rst = func(*args, **kwargs)
|
|
return rst
|
|
except AssertionError as e:
|
|
warnings.warn(f'{e}')
|
|
|
|
return wrapper
|