diff --git a/paddle/fluid/operators/uniform_random_op.cc b/paddle/fluid/operators/uniform_random_op.cc index 87699362b2b5a14750a01345098ec5e6cc9be115..a50add9739d8d21b2fc9f3de028553b12db01dfe 100644 --- a/paddle/fluid/operators/uniform_random_op.cc +++ b/paddle/fluid/operators/uniform_random_op.cc @@ -24,7 +24,15 @@ template class CPUUniformRandomKernel : public framework::OpKernel { public: void Compute(const framework::ExecutionContext& ctx) const override { - auto* tensor = ctx.Output("Out"); + framework::Tensor* tensor(nullptr); + auto out_var = ctx.OutputVar("Out"); + if (out_var->IsType()) { + tensor = ctx.Output("Out"); + } else if (out_var->IsType()) { + tensor = ctx.Output("Out")->mutable_value(); + } else { + PADDLE_THROW("Only support LoDTensor and SelectedRows."); + } T* data = tensor->mutable_data(ctx.GetPlace()); unsigned int seed = static_cast(ctx.Attr("seed")); std::minstd_rand engine; @@ -36,6 +44,7 @@ class CPUUniformRandomKernel : public framework::OpKernel { static_cast(ctx.Attr("min")), static_cast(ctx.Attr("max"))); int64_t size = tensor->numel(); + VLOG(3) << "size = " << size; for (int64_t i = 0; i < size; ++i) { data[i] = dist(engine); } @@ -55,6 +64,7 @@ class UniformRandomOp : public framework::OperatorWithKernel { "uniform_random's min must less then max"); auto& shape = ctx->Attrs().Get>("shape"); std::vector temp; + VLOG(3) << "shape.size() = " << shape.size(); temp.reserve(shape.size()); for (auto dim : shape) { temp.push_back(static_cast(dim));