Update trt_bert_creator.h

Tencent · Aug 17, 2021 · f247706 · f247706
1 parent d22837e
commit f247706
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/source/trt_engine/trt_network_crt/layer_creators/trt_bert_creator.h b/source/trt_engine/trt_network_crt/layer_creators/trt_bert_creator.h
@@ -214,6 +214,7 @@ class TLayerCreator<TrtBertDesc> : public ILayerCreator {
     if (bert_desc->hidden_size != 768 && bert_desc->hidden_size != 1024) use_int8 = false;
 
     const auto dtype = TrtCommon::GetDataType(bert_desc->use_fp16, use_int8, bert_desc->calib_mode);
+    const int has_skip = 1;
 
     nvinfer1::IPluginCreator* creator = getPluginRegistry()->getPluginCreator(
         fwd::bert::FWD_SKIP_LAYER_NORM_NAME, fwd::bert::FWD_SKIP_LAYER_NORM_VERSION);
@@ -223,6 +224,7 @@ class TLayerCreator<TrtBertDesc> : public ILayerCreator {
                             gamma.Count());
     field_data.emplace_back("ld", &bert_desc->hidden_size, nvinfer1::PluginFieldType::kINT32, 1);
     field_data.emplace_back("type_id", &dtype, nvinfer1::PluginFieldType::kINT32, 1);
+    field_data.emplace_back("has_skip", &has_skip, nvinfer1::PluginFieldType::kINT32, 1);
 
     const nvinfer1::PluginFieldCollection plugin_data{static_cast<int>(field_data.size()),
                                                       field_data.data()};