fix CI

nv-morpheus · Aug 29, 2024 · f4be468 · f4be468
1 parent 8033055
commit f4be468
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/python/morpheus/morpheus/_lib/src/messages/control.cpp b/python/morpheus/morpheus/_lib/src/messages/control.cpp
@@ -360,7 +360,7 @@ void ControlMessageProxy::payload_from_python_meta(ControlMessage& self, const p
 
 void ControlMessageProxy::set_tensors_from_python(ControlMessage& self, const pybind11::object& tensor_memory)
 {
-    TensorIndex count        = tensor_memory.attr("count").cast<TensorIndex>();
+    auto count               = tensor_memory.attr("count").cast<TensorIndex>();
     pybind11::object tensors = tensor_memory.attr("get_tensors")();
     self.tensors(TensorMemoryInterfaceProxy::init(count, tensors));
 }

diff --git a/tests/_utils/inference_worker.py b/tests/_utils/inference_worker.py
@@ -15,7 +15,7 @@
 
 import typing
 
-from morpheus.messages import MultiInferenceMessage
+from morpheus.messages import ControlMessage
 from morpheus.messages.memory.tensor_memory import TensorMemory
 from morpheus.stages.inference import inference_stage
 
@@ -30,5 +30,5 @@ def calc_output_dims(self, _):
         super().calc_output_dims(_)
         return (1, 2)
 
-    def process(self, _: MultiInferenceMessage, __: typing.Callable[[TensorMemory], None]):
+    def process(self, _: ControlMessage, __: typing.Callable[[TensorMemory], None]):
         raise NotImplementedError