The following code is used to specify device on which tf node is running
with tf.device('/gpu:0'):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
I have already known: this post, tensorflow doc and xla demo
What i want to know is:
Is there any way to specify XLA_GPU as the device on which tf node is running?
with tf.device('/XLA_GPU:0'):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
executing the code above gives ValueError: Unknown attribute: 'device' in '/job:localhost/replica:0/task:0/device:/XLA_GPU:0'
this is 100% reproducible on google colab.
/gpu:0
orgpu:0
. $\endgroup$with tf.device("device:XLA_GPU:0")
$\endgroup$