2
$\begingroup$

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.

$\endgroup$
2
  • $\begingroup$ Try /gpu:0 or gpu:0. $\endgroup$ – Esmailian Mar 16 '19 at 10:41
  • $\begingroup$ try with with tf.device("device:XLA_GPU:0") $\endgroup$ – Kiritee Gak Mar 16 '19 at 13:06
2
$\begingroup$

Check that the device is there first by examining the device names in the output from the following:

from tensorflow.python.client import device_lib
device_lib.list_local_devices()

Once you've verified that 'XLA_GPU:0' is a device on your system use something like

with tf.device('/device:XLA_GPU:0'):

Obviously you can change the names.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.