from inference_engine cimport InferenceEngine


cdef class TensorRTEngine(InferenceEngine):

    cdef public object context

    cdef public object d_input
    cdef public object d_output
    cdef str input_name
    cdef object input_shape

    cdef object h_output
    cdef str output_name
    cdef object output_shape

    cdef object stream


    cdef tuple get_input_shape(self)

    cdef int get_batch_size(self)

    cdef run(self, input_data)