如何在数据集上获得tensorflow日志设备位置？

import tensorflow as tf
from tensorflow.python.client import device_lib

local_devices = device_lib.list_local_devices()

gpu_devices = [device.name for device in local_devices if device.device_type == 'GPU']

dataset = tf.data.Dataset.from_tensor_slices(...)

dataset = dataset.interleave(lambda x: tf.data.Dataset.from_tensor_slices(x).with_options(device=tf.DeviceSpec(device_type="GPU", device_index=0)), cycle_length=len(gpu_devices), block_length=1)