[TAG] Fix signature of decorated function

python/setup.py

@@ -67,6 +67,7 @@ def config_cython():
     description='A domain specific language(DSL) for tensor computations.',
     install_requires=[
         'numpy',
+        'decorator',
         ],
     zip_safe=False,
     packages=[

python/tvm/tag.py

@@ -1,5 +1,5 @@
 """Tag class for TVM operators."""
-from functools import wraps
+from decorator import decorate
 
 class TagScope(object):
     """Tag scope object to set tag for operators, working as context
@@ -22,11 +22,10 @@ def __exit__(self, ptype, value, trace):
         TagScope.current = self._old_scope
 
     def __call__(self, fdecl):
-        @wraps(fdecl)
-        def tagged_fdecl(*args, **kwargs):
+        def tagged_fdecl(func, *args, **kwargs):
             with self:
-                return fdecl(*args, **kwargs)
-        return tagged_fdecl
+                return func(*args, **kwargs)
+        return decorate(fdecl, tagged_fdecl)
 
 
 def tag_scope(tag):

tests/ci_build/install/ubuntu_install_python.sh

@@ -4,5 +4,5 @@ apt-get update && apt-get install -y python-pip python-dev python3-dev
 # the version of the pip shipped with ubuntu may be too lower, install a recent version here
 cd /tmp && wget https://bootstrap.pypa.io/get-pip.py && python3 get-pip.py && python2 get-pip.py
 
-pip2 install nose pylint numpy nose-timer cython
-pip3 install nose pylint numpy nose-timer cython
+pip2 install nose pylint numpy nose-timer cython decorator
+pip3 install nose pylint numpy nose-timer cython decorator