diff --git a/tensorflow/python/util/tf_decorator.py b/tensorflow/python/util/tf_decorator.py
index a5d979e..8b4d22f 100644
--- a/tensorflow/python/util/tf_decorator.py
+++ b/tensorflow/python/util/tf_decorator.py
@@ -88,6 +88,7 @@
                           decorator_argspec)
   setattr(decorator_func, '_tf_decorator', decorator)
   decorator_func.__name__ = target.__name__
+  decorator_func.__module__ = target.__module__
   decorator_func.__doc__ = decorator.__doc__
   decorator_func.__wrapped__ = target
   return decorator_func