From c8e125872c2cffeb2018d8a98504ce15989a3b2d Mon Sep 17 00:00:00 2001
From: Li Fuchen <lfchener@outlook.com>
Date: Fri, 27 Sep 2019 11:00:34 +0800
Subject: [PATCH] Fixed warpctc, test=develop (#20011)

Use AllocateTmpTensor() for creating temporary tensors in warpctc.
---
 paddle/fluid/operators/warpctc_op.h | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/paddle/fluid/operators/warpctc_op.h b/paddle/fluid/operators/warpctc_op.h
index 1859c748d7..98c31e30dc 100644
--- a/paddle/fluid/operators/warpctc_op.h
+++ b/paddle/fluid/operators/warpctc_op.h
@@ -73,11 +73,12 @@ class WarpCTCFunctor {
                       "Bytes of workspace got by warp-ctc function, "
                       "get_workspace_size(), should be larger than 0.");
 
-    Tensor workspace;
+    auto& dev_ctx = ctx.template device_context<DeviceContext>();
     size_t workspace_elements = workspace_bytes / sizeof(float) + 1UL;
-    float* workspace_data = workspace.mutable_data<float>(
+    Tensor workspace = ctx.AllocateTmpTensor<float, DeviceContext>(
         framework::make_ddim({static_cast<int64_t>(workspace_elements)}),
-        ctx.GetPlace());
+        dev_ctx);
+    float* workspace_data = workspace.data<float>();
     math::SetConstant<DeviceContext, float>()(
         ctx.template device_context<DeviceContext>(), &workspace,
         static_cast<float>(0));
@@ -186,8 +187,10 @@ class WarpCTCKernel : public framework::OpKernel<T> {
         framework::make_ddim({static_cast<int64_t>(max_sequence_length),
                               static_cast<int64_t>(num_sequences),
                               static_cast<int64_t>(sequence_width)});
-    warpctc_logits.mutable_data<T>(warpctc_logits_dims, ctx.GetPlace());
-
+    auto& dev_ctx = ctx.template device_context<DeviceContext>();
+    Tensor warpctc_logits_tmp =
+        ctx.AllocateTmpTensor<T, DeviceContext>(warpctc_logits_dims, dev_ctx);
+    warpctc_logits.ShareDataWith(warpctc_logits_tmp);
     if (ctx.HasInput("LogitsLength")) {
       TensorCopySync(*logits, ctx.GetPlace(), &warpctc_logits);
     } else {
-- 
GitLab