Function allocate_binding_buffer

demo/HuggingFace/NNDF/tensorrt_utils.py:59–66

Allocate binding buffers for trt based on provided types and shapes dict

(types_dict, shapes_dict)

    return user_input

def allocate_binding_buffer(types_dict, shapes_dict):
    '''
    Allocate binding buffers for trt based on provided types and shapes dict
    '''
    return {
        k: torch.zeros(reduce(lambda v, a: v*a, shape), dtype=types_dict[k]).cuda()
        for k, shape in shapes_dict.items()
    }


def set_kv_data(kv_dict, past_or_present, layer_id, segment_value_dict):