tensorflow：有没有办法用tensorflow指定XLA_GPU？

Question

以下代码用于指定运行tf节点的设备

int foo = 10.9.floor(); // foo will be 10. 
int foo = 10.1.floor(); // foo will be 10. 

我已经知道： this post，tensorflow doc和xla demo

我想知道的是：

是否可以将XLA_GPU指定为运行tf节点的设备

with tf.device('/cpu:0'):
  a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')

执行上面的代码将得到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')

这在Google colab上是100％可重复的。

Answer 1

尝试'device：XLA_GPU：0'

1916                 SER TWAROGOWY WIEJSKI LUZ