class_docstr: Manage tensors stored in a compressed format.
class_docstr_tok: ["Manage","tensors","stored","in","a","compressed","format","."]
code: def init_cache_one_gpu_batch(self, config, task, policy):
      num_head, hidden_size, prompt_len, gen_len, gpu_batch_size = (
          config.n_head,
          config.input_dim,
          task.prompt_len,
          task.gen_len,
          policy.gpu_batch_size,
      )
      shape = (
          prompt_len + gen_len - 1,
          gpu_batch_size * num_head,
          hidden_size // num_head,
      )
      # NOTE: disable pin_memory due to high memory overhead
      pin_memory = False
      k_cache = self.allocate(
          shape,
          np.float16,
          comp_config=policy.comp_cache_config,
          pin_memory=pin_memory,
      )
      v_cache = self.allocate(
          shape,
          np.float16,
          comp_config=policy.comp_cache_config,
          pin_memory=pin_memory,
      )
      return k_cache, v_cache
  
imports: import dataclasses
  import torch
  import numpy
  from flexgen.pytorch_backend import TorchTensor, TorchDevice, DeviceType, general_copy, fix_recursive_import
  from flexgen.utils import np_dtype_to_torch_dtype
label_desc: Use image node_id 3 for calling the TorchCompressedDevice obj's underlying member method code with example usage: obj.init_cache_one_gpu_batch(config, task, policy) and returns: k_cache,  v_cache