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


    cpdef tuple get_input_shape(self)

    cpdef int get_batch_size(self)

    cpdef run(self, input_data)