diff --git a/paddle/fluid/operators/uniform_random_op_npu.cc b/paddle/fluid/operators/uniform_random_op_npu.cc index efd9d844fcb40b3714d55311c600262a0db868a5..580c1f3e948cc0bbd4041e2958ab2dcc2158a314 100644 --- a/paddle/fluid/operators/uniform_random_op_npu.cc +++ b/paddle/fluid/operators/uniform_random_op_npu.cc @@ -56,10 +56,13 @@ class NPUUniformRandomKernel : public framework::OpKernel { "unsupport type: %s.", framework::ToTypeName(out_var->Type()))); } - T *data = tensor->mutable_data(ctx.GetPlace()); - + tensor->mutable_data(ctx.GetPlace()); int64_t size = tensor->numel(); - std::unique_ptr data_cpu(new T[size]); + + Tensor cpu_tensor(tensor->type()); + cpu_tensor.Resize(tensor->dims()); + T *data_cpu = cpu_tensor.mutable_data(platform::CPUPlace()); + std::uniform_real_distribution dist( static_cast(ctx.Attr("min")), static_cast(ctx.Attr("max"))); @@ -90,12 +93,10 @@ class NPUUniformRandomKernel : public framework::OpKernel { } // copy to NPU - auto stream = - ctx.template device_context() - .stream(); - memory::Copy(BOOST_GET_CONST(platform::NPUPlace, ctx.GetPlace()), data, - platform::CPUPlace(), reinterpret_cast(data_cpu.get()), - size * sizeof(T), stream); + framework::TensorCopy( + cpu_tensor, ctx.GetPlace(), + ctx.template device_context(), tensor); + ctx.template device_context().Wait(); } }; diff --git a/python/paddle/fluid/tests/unittests/npu/test_uniform_random_op_npu.py b/python/paddle/fluid/tests/unittests/npu/test_uniform_random_op_npu.py index 8c37f0a32ac801112476547805cd2ceeb50f2aef..7c358c244f34ddf1c9b3d9e3f055ea333f9b435f 100644 --- a/python/paddle/fluid/tests/unittests/npu/test_uniform_random_op_npu.py +++ b/python/paddle/fluid/tests/unittests/npu/test_uniform_random_op_npu.py @@ -67,7 +67,7 @@ class TestNPUUniformRandomOp(OpTest): self.dtype = np.float32 def test_check_output(self): - self.check_output_customized(self.verify_output) + self.check_output_customized(self.verify_output, self.place) def verify_output(self, outs): hist, prob = self.output_hist(np.array(outs[0])) diff --git a/python/paddle/fluid/tests/unittests/op_test.py b/python/paddle/fluid/tests/unittests/op_test.py index f6de13b6fd4ce5980b930c1df06cb6676603ee45..02f1dfbb9f29e484ab6c661494f32a7bbcb29e21 100644 --- a/python/paddle/fluid/tests/unittests/op_test.py +++ b/python/paddle/fluid/tests/unittests/op_test.py @@ -1357,8 +1357,10 @@ class OpTest(unittest.TestCase): if self.op_type not in compile_vs_runtime_white_list.COMPILE_RUN_OP_WHITE_LIST: self.check_compile_vs_runtime(fetch_list, outs) - def check_output_customized(self, checker): + def check_output_customized(self, checker, custom_place=None): places = self._get_places() + if custom_place: + places.append(custom_place) for place in places: outs = self.calc_output(place) outs = [np.array(out) for out in outs]