code: def allocate_buffers(engine):
      """
      Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
      NOTE: currently this function only supports NetworkDefinitionCreationFlag::kEXPLICIT_BATCH flag.
  
      Parameters
      ----------
      engine : tensorrt.ICudaEngine
          An ICudaEngine for executing inference on a built network
  
      Returns
      -------
      list
          All input HostDeviceMem of an engine
      list
          All output HostDeviceMem of an engine
      GPU bindings
          Device bindings
      GPU stream
          A stream is a sequence of commands (possibly issued by different host threads) that execute in order
      """
      inputs = []
      outputs = []
      bindings = []
      stream = cuda.Stream()
      for binding in engine:
          size = trt.volume(
              engine.get_binding_shape(binding)
          )  # * engine.max_batch_size, batch size already in
          dtype = trt.nptype(engine.get_binding_dtype(binding))
          # Allocate host and device buffers
          host_mem = cuda.pagelocked_empty(size, dtype)
          device_mem = cuda.mem_alloc(host_mem.nbytes)
          # Append the device buffer to device bindings.
          bindings.append(int(device_mem))
          # Append to the appropriate list.
          if engine.binding_is_input(binding):
              inputs.append(HostDeviceMem(host_mem, device_mem))
          else:
              outputs.append(HostDeviceMem(host_mem, device_mem))
      return inputs, outputs, bindings, stream
  
file_path: nni/nni/compression/quantization_speedup/trt_pycuda.py
imports: import pycuda.driver
  import pycuda.autoinit
  import tensorrt
label_desc: Use image node_id 3 for calling a global function with example usage: allocate_buffers(engine) and returns: inputs,  outputs,  bindings,  stream