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, 2019 at 10:41
  • $\begingroup$ try with with tf.device("device:XLA_GPU:0") $\endgroup$ Mar 16, 2019 at 13:06

1 Answer 1

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 and acknowledge that you have read and understand our privacy policy and code of conduct.

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