Problem when testing "dropout" operator on GPU

Hello,
I write a simple example to test the “dropout” operator. At first, the example can work well on CPU. When I run the same example on GPU, it cannot be suppoted.
And it seems that the relative function SampleUniform can only be supported on CPU? When I remove the judgement sentence
if (data->ctx.device_type == kDLCPU)
I get the following error.
Segmentation fault: 11

Stack trace returned 10 entries:
[bt] (0) /usr/local/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x1e945a) [0x7fa094cdb45a]
[bt] (1) /usr/local/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x2f891f6) [0x7fa097a7b1f6]
[bt] (2) /lib/x86_64-linux-gnu/libc.so.6(+0x3ef20) [0x7fa172ff5f20]
[bt] (3) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(float* std::generate_n<float*, long, tvm::contrib::RandomEngine::SampleUniform(DLTensor*, float, float)::{lambda()#2}>(tvm::contrib::RandomEngine::SampleUniform(DLTensor*, float, float)::{lambda()#2}, long, tvm::contrib::RandomEngine::SampleUniform(DLTensor*, float, float)::{lambda()#2})+0x30) [0x7fa065ff53bc]
[bt] (4) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(tvm::contrib::RandomEngine::SampleUniform(DLTensor*, float, float)+0x3b4) [0x7fa065ff4aec]
[bt] (5) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(+0x25b755f) [0x7fa065ff355f]
[bt] (6) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(+0x25b7c38) [0x7fa065ff3c38]
[bt] (7) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(std::function<void (tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)>::operator()(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) const+0x5a) [0x7fa065491cb2]
[bt] (8) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(tvm::runtime::PackedFunc::CallPacked(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) const+0x30) [0x7fa0655fa8bc]
[bt] (9) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(TVMFuncCall+0x95) [0x7fa065f6d27a]
terminate called after throwing an instance of ‘dmlc::Error’
what(): [00:56:14] /home/yutong/git1009/DHCP_Compiler/tvm/src/runtime/workspace_pool.cc:116: Check failed: allocated_.size() == 1 (2 vs. 1) :
Stack trace:
[bt] (0) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(dmlc::LogMessageFatal::~LogMessageFatal()+0x34) [0x7fa06541491c]
[bt] (1) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(tvm::runtime::WorkspacePool::Pool::Release(DLContext, tvm::runtime::DeviceAPI*)+0xff) [0x7fa065f909f5]
[bt] (2) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(tvm::runtime::WorkspacePool::~WorkspacePool()+0xab) [0x7fa065f8fdd3]
[bt] (3) /home/yutong/git1009/DHCP_Compiler/tvm/build/libtvm.so(tvm::runtime::CUDAThreadEntry::~CUDAThreadEntry()+0x1c) [0x7fa065fe75b0]
[bt] (4) /lib/x86_64-linux-gnu/libc.so.6(__call_tls_dtors+0x3f) [0x7fa172ffa8af]
[bt] (5) /lib/x86_64-linux-gnu/libc.so.6(+0x43117) [0x7fa172ffa117]
[bt] (6) /lib/x86_64-linux-gnu/libc.so.6(+0x4313a) [0x7fa172ffa13a]
[bt] (7) /usr/local/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x2f8922e) [0x7fa097a7b22e]
[bt] (8) /lib/x86_64-linux-gnu/libc.so.6(+0x3ef20) [0x7fa172ff5f20]

Any advice on supporting “dropout” on GPU? Thanks a lot!:grinning: