tf.device

Устанавливает устройство для операций, созданных/выполненных в этом контексте.

tf.device(
    device_name
)

Эта функция определяет устройство, которое будет использоваться для операций, созданных/выполненных в определенном контексте. Вложенные контексты будут наследовать и также создавать/выполнять свои операции на указанном устройстве. Если конкретное устройство не требуется, следует воздержаться от использования этой функции, чтобы устройство можно было автоматически назначить. В целом использование этой функции необязательно. device_name может быть полностью указано, как в "/job:worker/task:1/device:cpu:0", или частично указано, содержа в себе только подмножество полей, разделенных символом "/". Любые указанные поля переопределят аннотации устройства из внешних областей.

Например:

with tf.device('/job:foo'):
  # ops created here have devices with /job:foo
  with tf.device('/job:bar/task:0/device:gpu:2'):
    # ops created here have the fully specified device above
  with tf.device('/device:gpu:1'):
    # ops created here have the device '/job:foo/device:gpu:1'