TensorFlow GPU, CUDA_ERROR_LAUNCH_FAILED on tf.one_hot()

Question

I try to run tf.one_hot, get CUDA_ERROR_LAUNCH_FAILED error. Here is the details:

Sample code:

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)

Expected result:

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

Actual result:

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

Config of the PC:

• TensorFlow 0.12.0-rc1
• Python 3.5
• CUDA 8.0
• cuDNN 5.1
• OS: Windows 10
• GPU: GeForce GTX 970

tf.one_hot run ok when running on Linux CPU, Linux GPU (GeForce GTX 660), Windows 10 CPU. Not ok on the Windows 10 GPU.

On the Windows 10 GPU, tf.matmul, tf.reduce_mean, tf.reduce_sum are run ok. But tf.one_hot is not ok.

Is that a bug, or I miss something? Thanks.

(Edit 2016-12-16)

I have run the code on the same machine, in Xubuntu, GPU. The code run fine. So I think that is a problem in TensorFlow-Windows.

Answer 1

Also leaving a comment as an answer, in my case because I don't have enough reputation to comment.

Have you reported this as a bug on GitHub? I too can confirm this behaviour with the same Tensorflow/OS/Graphics Card/etc

Just moving the tf.one_hot() to the CPU solves the issue for me, i.e. something like

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