From aa7f9c7167edc16a2d81c09971d6d0fcdc989aa1 Mon Sep 17 00:00:00 2001
From: wxl <wanxuelei@huawei.com>
Date: Wed, 31 Mar 2021 10:42:25 +0800
Subject: [PATCH] fix 1951 ts 4g memory failed

---
 ge/graph/load/model_manager/davinci_model.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/ge/graph/load/model_manager/davinci_model.cc b/ge/graph/load/model_manager/davinci_model.cc
index 6dc3c0f3..cd78e5b8 100755
--- a/ge/graph/load/model_manager/davinci_model.cc
+++ b/ge/graph/load/model_manager/davinci_model.cc
@@ -3094,12 +3094,13 @@ Status DavinciModel::MallocKnownArgs() {
       return ret;
     }
   }
+  rtError_t rt_ret;
   bool is_support = false;
   GE_CHK_STATUS_RET_NOLOG(CheckCapability(FEATURE_TYPE_MEMORY, MEMORY_INFO_TS_4G_LIMITED, is_support));
   auto mem_type = is_support ? RT_MEMORY_TS_4G : RT_MEMORY_HBM;
   // malloc args memory
   if (total_args_size_ != 0) {
-    rt_ret = rtMalloc(&args_, total_args_size_, RT_MEMORY_HBM);
+    rt_ret = rtMalloc(&args_, total_args_size_, mem_type);
     if (rt_ret != RT_ERROR_NONE) {
       REPORT_CALL_ERROR("E19999", "Call rtMalloc failed, size:%u, ret: 0x%X when DavinciModel %s",
                         total_args_size_, rt_ret, __FUNCTION__);