diff --git a/paddle/fluid/operators/linspace_op.cc b/paddle/fluid/operators/linspace_op.cc index 667c6e892956e29478f1401c3cb2622713433037..7cc07383bfa5f67a2404b220cb481d9017b40fd8 100644 --- a/paddle/fluid/operators/linspace_op.cc +++ b/paddle/fluid/operators/linspace_op.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. */ #include "paddle/fluid/operators/linspace_op.h" +#include namespace paddle { namespace operators { @@ -21,7 +22,7 @@ class LinspaceOp : public framework::OperatorWithKernel { public: using framework::OperatorWithKernel::OperatorWithKernel; - void InferShape(framework::InferShapeContext* ctx) const override { + void InferShape(framework::InferShapeContext *ctx) const override { OP_INOUT_CHECK(ctx->HasInput("Start"), "Input", "Start", "linspace"); OP_INOUT_CHECK(ctx->HasInput("Stop"), "Input", "Stop", "linspace"); OP_INOUT_CHECK(ctx->HasInput("Num"), "Input", "Num", "linspace"); @@ -50,11 +51,17 @@ class LinspaceOp : public framework::OperatorWithKernel { protected: framework::OpKernelType GetExpectedKernelType( - const framework::ExecutionContext& ctx) const override { + const framework::ExecutionContext &ctx) const override { return framework::OpKernelType( framework::proto::VarType::Type(ctx.Attr("dtype")), ctx.GetPlace()); } + + framework::OpKernelType GetKernelTypeForVar( + const std::string &var_name, const framework::Tensor &tensor, + const framework::OpKernelType &expected_kernel_type) const override { + return expected_kernel_type; + } }; class LinspaceOpMaker : public framework::OpProtoAndCheckerMaker { diff --git a/python/paddle/fluid/layers/tensor.py b/python/paddle/fluid/layers/tensor.py index cf52f3b00fb2739d186021dc51d6aa0f506be706..2fba578ec077f2a74388d433bf3ab5b3098e81ad 100644 --- a/python/paddle/fluid/layers/tensor.py +++ b/python/paddle/fluid/layers/tensor.py @@ -1453,11 +1453,14 @@ def linspace(start, stop, num, dtype=None, name=None): if not isinstance(dtype, core.VarDesc.VarType): dtype = convert_np_dtype_to_dtype_(dtype) if not isinstance(start, Variable): - tensor_start = fill_constant([1], dtype, start) + with device_guard("cpu"): + tensor_start = fill_constant([1], dtype, start) if not isinstance(stop, Variable): - tensor_stop = fill_constant([1], dtype, stop) + with device_guard("cpu"): + tensor_stop = fill_constant([1], dtype, stop) if not isinstance(num, Variable): - tensor_num = fill_constant([1], 'int32', num) + with device_guard("cpu"): + tensor_num = fill_constant([1], 'int32', num) if in_dygraph_mode(): return core.ops.linspace(tensor_start, tensor_stop, tensor_num, 'dtype', dtype)