dmlc / MXNet.jl

MXNet Julia Package - flexible and efficient deep learning in Julia

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Small bug makes the kernel crash when using softmax_cross_entropy on executor outputs on gpu.

albinsunesson opened this issue · comments

If you calculate the cross entropy on a NDArray from an executor output like this:

mx.softmax_cross_entropy(exec.outputs[1],label)

The kernel crashes.

Here is a small code example to reproduce the error:
data = mx.Variable(:data)
fc1 = mx.FullyConnected(data, num_hidden=128, name=:fc1)
act1 = mx.Activation(fc1, name=:act1, act_type="tanh")
fc2 = mx.FullyConnected(act1, name=:fc2, num_hidden=10)
softmax = mx.SoftmaxOutput(fc2, name=:softmax)

batch_size = 100
data_shape = (28*28,batch_size)
exec = mx.simple_bind(softmax, mx.gpu() , data=data_shape)

label = mx.ones(100)
mx.softmax_cross_entropy(exec.outputs[1],label)

However, if you create a new NDArray and calculate the cross entropy like this everything works fine.

pred = mx.zeros(size(exec.outputs[1]))
pred[:] = exec.outputs[1]
mx.softmax_cross_entropy(pred,label)

Is that because the two NDArray lives on different devices? (CPU and GPU).