Function decoder_inference

demo/HuggingFace/T5/measurements.py:33–50

(
    t5_decoder, input_ids, encoder_last_hidden_state, timing_profile, use_cuda=True, use_cache=False, past_key_values=None
)

@use_cuda
def decoder_inference(
    t5_decoder, input_ids, encoder_last_hidden_state, timing_profile, use_cuda=True, use_cache=False, past_key_values=None
):
    # This implementation is a bit ugly. Moving implementation of the model to check HFRunner would be cleaner.
    if isinstance(t5_decoder, TRTNativeRunner):
        # Function is technically in T5TRTDecoder however due to circular import, TRTNativeRunner in this module scope
        # implies the existence of this function.
        t5_decoder.set_return_device("cuda" if use_cuda else "cpu")

    def decoder_stmt():
        t5_decoder(
            input_ids=input_ids, encoder_hidden_states=encoder_last_hidden_state, use_cache=use_cache,
            past_key_values=past_key_values
        )

    decoder_e2e_time = measure_python_inference_code(decoder_stmt, timing_profile)

    return (decoder_stmt(), decoder_e2e_time)


@use_cuda