TensorFlow GPU，tf.one_hot() 上的 CUDA_ERROR_LAUNCH_FAILED

Question

我尝试运行 tf.one_hot，出现 CUDA_ERROR_LAUNCH_FAILED 错误。详情如下:

示例代码:

import tensorflow as tf
idx_0 = tf.placeholder(tf.int64, [None])
mask = tf.one_hot(idx_0, 3, axis=-1)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
a = sess.run([mask],feed_dict={idx_0:[0,1,2]})
print(a)

预期结果:

[array([[ 1.,  0.,  0.],
       [ 0.,  1.,  0.],
       [ 0.,  0.,  1.]], dtype=float32)]

实际结果:

E c:\tf_jenkins\home\workspace\release-win\device\gpu\os\windows\tensorflow\stream_executor\cuda\cuda_driver.cc:1177] could not synchronize on CUDA context: CUDA_ERROR_LAUNCH_FAILED :: No stack trace available
E c:\tf_jenkins\home\workspace\release-win\device\gpu\os\windows\tensorflow\stream_executor\cuda\cuda_event.cc:49] Error polling for event status: failed to query event: CUDA_ERROR_LAUNCH_FAILED
F c:\tf_jenkins\home\workspace\release-win\device\gpu\os\windows\tensorflow\core\common_runtime\gpu\gpu_util.cc:370] GPU sync failed

电脑配置:

• TensorFlow 0.12.0-rc1
• python 3.5
• CUDA 8.0
• cuDNN 5.1
• 操作系统:Windows 10
• 图形处理器:GeForce GTX 970

tf.one_hot 在 Linux CPU、Linux GPU (GeForce GTX 660)、Windows 10 CPU 上运行时运行正常。在 Windows 10 GPU 上不正常。

在 Windows 10 GPU 上，tf.matmul、tf.reduce_mean、tf.reduce_sum 运行正常。但是 tf.one_hot 不行。

这是一个错误，还是我遗漏了什么？谢谢。

(编辑 2016-12-16)

我在同一台机器上的 Xubuntu 和 GPU 中运行了代码。代码运行良好。所以我认为这是 TensorFlow-Windows 中的一个问题。

Answer 1

同时留下评论作为答案，在我的例子中是因为我没有足够的声誉来发表评论。

您是否将此报告为 GitHub 上的错误？我也可以使用相同的 Tensorflow/OS/Graphics Card/etc 确认此行为

只需将 tf.one_hot() 移至 CPU 即可为我解决问题，即类似于

与 tf.device('/cpu:0'):
b = tf.one_hot(a, 123)