Add compute_schema method to fix control message support

nv-morpheus · Oct 31, 2023 · e3f43c2 · e3f43c2
1 parent a66edc4
commit e3f43c2
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/morpheus/stages/preprocess/deserialize_stage.py b/morpheus/stages/preprocess/deserialize_stage.py
@@ -29,6 +29,7 @@
 from morpheus.messages import MessageMeta
 from morpheus.messages import MultiMessage
 from morpheus.pipeline.multi_message_stage import MultiMessageStage
+from morpheus.pipeline.stage_schema import StageSchema
 
 logger = logging.getLogger(__name__)
 
@@ -94,6 +95,9 @@ def supports_cpp_node(self):
         # Enable support by default
         return True
 
+    def compute_schema(self, schema: StageSchema):
+        schema.output_schema.set_type(self._message_type)
+
     @staticmethod
     def check_slicable_index(x: MessageMeta, ensure_sliceable_index: bool = True):
         if (not x.has_sliceable_index()):