code: def get_workload(
      batch_size,
      num_actions=18,
      image_shape=(4, 84, 84),
      dtype="float32",
      layout="NCHW",
  ):
      """Get benchmark workload for a Deep Q Network
      Parameters
      ----------
      batch_size : int
          The batch size used in the model
      num_actions : int, optional
          Number of actions
      image_shape : tuple, optional
          The input image shape
      dtype : str, optional
          The data type
      Returns
      -------
      mod : tvm.IRModule
          The relay module that contains a DQN network.
      params : dict of str to NDArray
          The parameters.
      """
      net = get_net(
          batch_size,
          num_actions=num_actions,
          image_shape=image_shape,
          dtype=dtype,
          layout=layout,
      )
      return create_workload(net)
  
imports: from tvm import relay
  from .None import layers
  from .init import create_workload
label_desc: Use image node_id 2 for calling a global function with example usage: get_workload(batch_size, num_actions, image_shape, dtype, layout) and returns: create_workload