Skip to content

Commit f12b564

Browse files
committed
merge cleanups
1 parent d6beeaf commit f12b564

File tree

2 files changed

+3
-3
lines changed

2 files changed

+3
-3
lines changed

python/morpheus/morpheus/messages/__init__.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
# Import order is very important here. Import base classes before child ones
1919
# isort: off
2020

21-
from morpheus._lib.messages import ControlMessage
2221
from morpheus._lib.messages import DataLoaderRegistry
2322
from morpheus._lib.messages import RawPacketMessage
2423
from morpheus.messages.memory.tensor_memory import TensorMemory
@@ -33,6 +32,7 @@
3332
from morpheus.messages.message_meta import MessageMeta
3433
from morpheus.messages.multi_message import MultiMessage
3534
from morpheus.messages.message_meta import UserMessageMeta
35+
from morpheus.messages.control_message import ControlMessage
3636

3737
__all__ = [
3838
"ControlMessage",

python/morpheus/morpheus/stages/preprocess/preprocess_ae_stage.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,11 @@
1818

1919
import cupy as cp
2020

21-
import morpheus._lib.messages as _messages
2221
from morpheus.cli.register_stage import register_stage
2322
from morpheus.config import Config
2423
from morpheus.config import PipelineModes
2524
from morpheus.messages import ControlMessage
25+
from morpheus.messages import TensorMemory
2626
from morpheus.stages.preprocess.preprocess_base_stage import PreprocessBaseStage
2727

2828
logger = logging.getLogger(__name__)
@@ -101,7 +101,7 @@ def pre_process_batch(msg: ControlMessage, fea_len: int, feature_columns: typing
101101
msg.set_metadata("model", autoencoder)
102102
msg.set_metadata("train_scores_mean", scores_mean)
103103
msg.set_metadata("train_scores_std", scores_std)
104-
msg.tensors(_messages.TensorMemory(count=count, tensors={"input": inputs, "seq_ids": seg_ids}))
104+
msg.tensors(TensorMemory(count=count, tensors={"input": inputs, "seq_ids": seg_ids}))
105105
return msg
106106

107107
def _get_preprocess_fn(self) -> typing.Callable[[ControlMessage], ControlMessage]:

0 commit comments

Comments
 (0)