Работает ли tf.contrib.eager.list_devices () на машинах AI Platform?

Question

Я пытаюсь настроить учебное задание Tensorflow на платформе AI и хочу иметь возможность динамически настраивать количество используемых графических процессоров.В настоящее время, чтобы получить количество графических процессоров, я использую:

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))

Однако, похоже, это работает только при локальном запуске.Почему это не работает, когда я запускаю свою работу на AI Platform с workerType complex_model_m_gpu, и какие альтернативы я могу использовать для достижения желаемого эффекта?

Дополнительная информация: мой файл конфигурации выглядит следующим образом:

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