code: def set_num_kv_heads(num):
      global num_kv_heads
      num_kv_heads = num
  
imports: from deepspeed import comm
label_desc: Use image node_id 1 for calling a global function with example usage: set_num_kv_heads(num) without return types