Method mark_output

tensorrt_llm/functional.py:299–317

Mark a tensor as a network output. When a tensor is marked as an output, its content can be obtained after the execution of the TensorRT engine. The user is responsible for allocating buffers to store the output tensors when preparing the execution of the Te

(self,
                    name: Optional[str] = None,
                    dtype: Optional[Union[str, trt.DataType]] = None)

            self.trt_tensor.location = location

    def mark_output(self,
                    name: Optional[str] = None,
                    dtype: Optional[Union[str, trt.DataType]] = None):
        '''
        Mark a tensor as a network output.

        When a tensor is marked as an output, its content can be obtained after
        the execution of the TensorRT engine. The user is responsible for
        allocating buffers to store the output tensors when preparing the
        execution of the TensorRT engine.
        '''
        if name is None:
            name = self.name

        if isinstance(dtype, str):
            dtype = str_dtype_to_trt(dtype)

        assert dtype is None or isinstance(dtype, trt.DataType)
        default_net()._mark_output(self, name, dtype)

    def __add__(self, b):
        '''