diff --git a/ge/common/helper/model_helper.cc b/ge/common/helper/model_helper.cc
index e552424c..454e6464 100644
--- a/ge/common/helper/model_helper.cc
+++ b/ge/common/helper/model_helper.cc
@@ -104,7 +104,7 @@ Status ModelHelper::SaveSizeToModelDef(const GeModelPtr &ge_model, ge::Buffer &m
   GELOGD("SaveSizeToModelDef task_info_size is %zu", partition_task_size);
   om_info.push_back(partition_task_size);
 
-  GE_CHK_BOOL_EXEC(ge::AttrUtils::SetListInt(ge_model, "om_info_list", om_info),
+  GE_CHK_BOOL_EXEC(ge::AttrUtils::SetListInt(*(ge_model.get()), "om_info_list", om_info),
                    GELOGE(FAILED, "SetListInt of om_info_list failed.");
                    return FAILED);
 
@@ -201,10 +201,6 @@ Status ModelHelper::SaveModelTaskDef(std::shared_ptr<OmFileSaveHelper> &om_file_
   GELOGD("TASK_INFO op_size:%d, stream_num:%u", model_task_def->op().size(), model_task_def->stream_num());
   GELOGD("TASK_INFO size is %zu", partition_task_size);
 
-  GE_CHK_BOOL_EXEC(ge::AttrUtils::SetInt(ge_model, "task_info_size", partition_task_size),
-                   GELOGE(FAILED, "SetInt of task_info_size failed.");
-                   return FAILED);
-
   if (SaveModelPartition(om_file_save_helper, ModelPartitionType::TASK_INFO, task_buffer.GetData(),
                          partition_task_size, model_index) != SUCCESS) {
     GELOGE(PARAM_INVALID, "Add model task def partition failed");