diff --git a/tensorflow/virtual_root_template_v1.__init__.py b/tensorflow/virtual_root_template_v1.__init__.py index e3049cee645..236e9f52258 100644 --- a/tensorflow/virtual_root_template_v1.__init__.py +++ b/tensorflow/virtual_root_template_v1.__init__.py @@ -54,6 +54,9 @@ class _LazyLoader(_types.ModuleType): module = self._load() return dir(module) + def __reduce__(self): + return __import__, (self.__name__,) + # Forwarding a module is as simple as lazy loading the module from the new path # and then registering it to sys.modules using the old path diff --git a/tensorflow/virtual_root_template_v2.__init__.py b/tensorflow/virtual_root_template_v2.__init__.py index e09127d70a6..83c020182a8 100644 --- a/tensorflow/virtual_root_template_v2.__init__.py +++ b/tensorflow/virtual_root_template_v2.__init__.py @@ -54,6 +54,9 @@ class _LazyLoader(_types.ModuleType): module = self._load() return dir(module) + def __reduce__(self): + return __import__, (self.__name__,) + # Forwarding a module is as simple as lazy loading the module from the new path # and then registering it to sys.modules using the old path