未验证 提交 21a72c55 编写于 作者: zhouweiwei2014's avatar zhouweiwei2014 提交者: GitHub

fix multinomial memory leak bug (#51851)

上级 65e3fa35
......@@ -145,18 +145,9 @@ void MultinomialKernel(const Context& dev_ctx,
int64_t in_data_numel = x.numel();
int64_t out_data_numel = out->numel();
// Just use to PADDLE_ENFORCE error message
T* cpu_in_data = new T[in_data_numel];
#ifdef PADDLE_WITH_HIP
hipMemcpy(
cpu_in_data, in_data, in_data_numel * sizeof(T), hipMemcpyDeviceToHost);
#else
cudaMemcpy(cpu_in_data,
in_data,
in_data_numel * sizeof(T),
cudaMemcpyDeviceToHost);
#endif
phi::DenseTensor cpu_tensor;
phi::Copy<Context>(dev_ctx, x, phi::CPUPlace(), false, &cpu_tensor);
T* cpu_in_data = cpu_tensor.data<T>();
for (size_t i = 0; i < num_distributions; ++i) {
int zero_num = 0;
for (size_t j = 0; j < num_categories; ++j) {
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册