diff --git a/python/morpheus/morpheus/messages/__init__.py b/python/morpheus/morpheus/messages/__init__.py
index a0e4e92953..7186b35fb7 100644
--- a/python/morpheus/morpheus/messages/__init__.py
+++ b/python/morpheus/morpheus/messages/__init__.py
@@ -18,7 +18,6 @@
 # Import order is very important here. Import base classes before child ones
 # isort: off
 
-from morpheus._lib.messages import ControlMessage
 from morpheus._lib.messages import DataLoaderRegistry
 from morpheus._lib.messages import RawPacketMessage
 from morpheus.messages.memory.tensor_memory import TensorMemory
@@ -33,6 +32,7 @@
 from morpheus.messages.message_meta import MessageMeta
 from morpheus.messages.multi_message import MultiMessage
 from morpheus.messages.message_meta import UserMessageMeta
+from morpheus.messages.control_message import ControlMessage
 
 __all__ = [
     "ControlMessage",
diff --git a/python/morpheus/morpheus/stages/preprocess/preprocess_ae_stage.py b/python/morpheus/morpheus/stages/preprocess/preprocess_ae_stage.py
index 6b29195682..9163d5491f 100644
--- a/python/morpheus/morpheus/stages/preprocess/preprocess_ae_stage.py
+++ b/python/morpheus/morpheus/stages/preprocess/preprocess_ae_stage.py
@@ -18,11 +18,11 @@
 
 import cupy as cp
 
-import morpheus._lib.messages as _messages
 from morpheus.cli.register_stage import register_stage
 from morpheus.config import Config
 from morpheus.config import PipelineModes
 from morpheus.messages import ControlMessage
+from morpheus.messages import TensorMemory
 from morpheus.stages.preprocess.preprocess_base_stage import PreprocessBaseStage
 
 logger = logging.getLogger(__name__)
@@ -101,7 +101,7 @@ def pre_process_batch(msg: ControlMessage, fea_len: int, feature_columns: typing
         msg.set_metadata("model", autoencoder)
         msg.set_metadata("train_scores_mean", scores_mean)
         msg.set_metadata("train_scores_std", scores_std)
-        msg.tensors(_messages.TensorMemory(count=count, tensors={"input": inputs, "seq_ids": seg_ids}))
+        msg.tensors(TensorMemory(count=count, tensors={"input": inputs, "seq_ids": seg_ids}))
         return msg
 
     def _get_preprocess_fn(self) -> typing.Callable[[ControlMessage], ControlMessage]: