Fix type-hint

nv-morpheus · Dec 20, 2023 · bee89d3 · bee89d3
1 parent 6ea799b
commit bee89d3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/examples/log_parsing/inference.py b/examples/log_parsing/inference.py
@@ -23,6 +23,7 @@
 from morpheus.config import Config
 from morpheus.config import PipelineModes
 from morpheus.messages import MultiInferenceMessage
+from morpheus.messages import MultiInferenceNLPMessage
 from morpheus.messages import MultiResponseMessage
 from morpheus.messages import TensorMemory
 from morpheus.pipeline.stage_schema import StageSchema
@@ -142,7 +143,7 @@ def compute_schema(self, schema: StageSchema):
         schema.output_schema.set_type(MultiResponseMessage)
 
     @staticmethod
-    def _convert_one_response(output: MultiResponseMessage, inf: MultiInferenceMessage,
+    def _convert_one_response(output: MultiResponseMessage, inf: MultiInferenceNLPMessage,
                               res: TensorMemory) -> MultiResponseMessage:
 
         output.input_ids[inf.offset:inf.count + inf.offset, :] = inf.input_ids