From 50f149a48e7f32b884ffbb88b0e04f91a21ecf7f Mon Sep 17 00:00:00 2001
From: Zhaolong Xing <nhzlx.dragon@gmail.com>
Date: Mon, 10 Aug 2020 13:25:12 +0800
Subject: [PATCH] fix cudnn workspace size problem during inference. (#26021)

test=develop
---
 paddle/fluid/operators/fused/conv_fusion_op.cu | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/paddle/fluid/operators/fused/conv_fusion_op.cu b/paddle/fluid/operators/fused/conv_fusion_op.cu
index e8f371cb487..b22f28fbbe3 100644
--- a/paddle/fluid/operators/fused/conv_fusion_op.cu
+++ b/paddle/fluid/operators/fused/conv_fusion_op.cu
@@ -216,6 +216,12 @@ class CUDNNConvFusionOpKernel : public framework::OpKernel<T> {
               perf_results.get()));
       algo = (perf_results.get())[best_algo_idx].algo;
       VLOG(3) << "cuDNN forward algo " << algo;
+      PADDLE_ENFORCE_CUDA_SUCCESS(
+          platform::dynload::cudnnGetConvolutionForwardWorkspaceSize(
+              handle, cudnn_input_desc, cudnn_filter_desc, cudnn_conv_desc,
+              cudnn_output_desc, algo, &workspace_size_in_bytes));
+      if (workspace_size_in_bytes > workspace_size_limit)
+        workspace_size_limit = workspace_size_in_bytes;
     } else {
       std::function<cudnnConvolutionFwdAlgo_t()> search_func =
           [&]() -> cudnnConvolutionFwdAlgo_t {
-- 
GitLab