diff --git a/src/opr/impl/misc.cpp b/src/opr/impl/misc.cpp index ed35b122dac4be6832f4cd7722af3f1334aed3c0..93931680793fabcf1f285e92e1f2d5f1f507a878 100644 --- a/src/opr/impl/misc.cpp +++ b/src/opr/impl/misc.cpp @@ -379,6 +379,8 @@ void TopK::init_output_static_infer_desc() { } auto infer_workspace = [this](TensorShape& dst, const InpVal& iv) { + // active comp_node for cuda launch kernel in get_workspace_in_bytes + comp_node().activate(); auto k = iv.val[3].value().ptr()[0]; auto size = megdnn_opr()->get_workspace_in_bytes( k, {iv.val[0].shape(), input(0)->dtype()}, diff --git a/src/opr/impl/rand.cpp b/src/opr/impl/rand.cpp index cbbcc64562598fd34f20070de7581952562d650e..e25d5f9167fbcfc98a18335af568739337372110 100644 --- a/src/opr/impl/rand.cpp +++ b/src/opr/impl/rand.cpp @@ -60,6 +60,8 @@ cg::OperatorNodeBase::NodeProp* RNGOprBase::do_make_node_prop() const { void RNGOprBase::ensure_megdnn_opr() { if (!m_megdnn_opr || m_megdnn_opr.comp_node() != comp_node()) { + // activate comp_node for curandCreateGenerator in create_megdnn_opr + comp_node().activate(); m_megdnn_opr = create_megdnn_opr(); } }