tf.contrib.eager.list_devices（）是否可以在AI Platform计算机上使用？

Question

我正在尝试在AI平台上设置Tensorflow培训工作，并希望能够动态配置要使用的GPU数量。目前，要获取GPU的数量，我正在使用：

distribution_strategy = None
# Get the available GPU devices
num_gpus = len([device_name
                for device_name in tf.contrib.eager.list_devices()
                if '/device:GPU' in device_name])
logging.info('%s GPUs are available.', str(num_gpus))
if num_gpus > 1:
  distribution_strategy = tf.distribute.MirroredStrategy()
  logging.info('MirroredStrategy will be used for training.')
  # Update the batch size
  args.batch_size = int(math.ceil(args.batch_size / num_gpus))

但是，这似乎仅在本地运行时有效。当我使用workerType complex_model_m_gpu在AI平台上运行我的工作时，为什么这不起作用？我可以使用哪些替代方法来达到预期的效果？

更多信息：我的配置文件如下：

trainingInput:
  scaleTier: CUSTOM
  masterType: large_model
  workerType: complex_model_m_gpu
  parameterServerType: large_model
  workerCount: 3
  parameterServerCount: 1