diff --git a/python/mxnet/__init__.py b/python/mxnet/__init__.py index e960829e691b..374a3b50bbb5 100644 --- a/python/mxnet/__init__.py +++ b/python/mxnet/__init__.py @@ -54,7 +54,6 @@ from . import lr_scheduler # use mx.kv as short for kvstore from . import kvstore as kv -from . import kvstore_server # Runtime compile module from . import rtc # Attribute scope to add attributes to symbolic graphs @@ -82,3 +81,11 @@ from . import gluon __version__ = base.__version__ + +# Dist kvstore module which launches a separate process when role is set to "server". +# This should be done after other modules are initialized. +# Otherwise this may result in errors when unpickling custom LR scheduler/optimizers. +# For example, the LRScheduler in gluoncv depends on a specific version of MXNet, and +# checks the __version__ attr of MXNet, which is not set on kvstore server due to the +# fact that kvstore-server module is imported before the __version__ attr is set. +from . import kvstore_server