code: def get_num_gpus(flags_obj):
      """Treat num_gpus=-1 as 'use all'."""
      if flags_obj.num_gpus != -1:
          return flags_obj.num_gpus
  
      from tensorflow.python.client import (
          device_lib,
      )  # pylint: disable=g-import-not-at-top
  
      local_device_protos = device_lib.list_local_devices()
      return sum(
          [1 for d in local_device_protos if d.device_type == "GPU"]
      )
  
file_path: autodist/examples/benchmark/utils/flags/_base.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl import flags
  import tensorflow
  from utils.flags._conventions import help_wrap
  from utils.logs import hooks_helper
label_desc: Use image node_id 2 for calling a global function with example usage: get_num_gpus(flags_obj) and returns: sum, flags_obj