[OP] Update BiasAdd OP settings

Tencent · Mar 23, 2024 · 1e2b6a6 · 1e2b6a6
1 parent 4d3dc00
commit 1e2b6a6
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/source/tnn/device/cpu/acc/cpu_bias_add_layer_acc.cc b/source/tnn/device/cpu/acc/cpu_bias_add_layer_acc.cc
@@ -25,9 +25,9 @@ Status CpuBiasAddLayerAcc::Reshape(const std::vector<Blob *> &inputs, const std:
 }
 
 Status CpuBiasAddLayerAcc::Forward(const std::vector<Blob *> &inputs, const std::vector<Blob *> &outputs) {
-    auto resource = dynamic_cast<BiasAddLayerResource *>(resource_);
+    auto resource = dynamic_cast<BatchNormLayerResource *>(resource_);
     if (!resource) {
-        return Status(TNNERR_MODEL_ERR, "Error: BiasAddLayerResource is nil");
+        return Status(TNNERR_MODEL_ERR, "Error: BatchNormLayerResource is nil");
     }
 
     auto input_blob        = inputs[0];

diff --git a/source/tnn/interpreter/tnn/layer_interpreter/bias_add_layer_interpreter.cc b/source/tnn/interpreter/tnn/layer_interpreter/bias_add_layer_interpreter.cc
@@ -33,7 +33,7 @@ Status BiasAddLayerInterpreter::SaveProto(std::ostream& output_stream, LayerPara
 }
 
 Status BiasAddLayerInterpreter::SaveResource(Serializer& serializer, LayerParam* param, LayerResource* resource) {
-    CAST_OR_RET_ERROR(bias_res, BiasAddLayerResource, "invalid layer res to save", resource);
+    CAST_OR_RET_ERROR(bias_res, BatchNormLayerResource, "invalid layer res to save", resource);
     serializer.PutRaw(bias_res->bias_handle);
     return TNN_OK;
 }

diff --git a/source/tnn/network/tensorrt/layer_builder/batch_norm_layer_builder.cc b/source/tnn/network/tensorrt/layer_builder/batch_norm_layer_builder.cc
@@ -75,6 +75,7 @@ ILayer* BatchNormTRTLayerBuilder::AddToNetwork(INetworkDefinition* network) {
 }
 
 REGISTER_TENSORRT_LAYER_BUILDER(BatchNorm, LAYER_BATCH_NORM);
+REGISTER_TENSORRT_LAYER_BUILDER(BatchNorm, LAYER_BIAS_ADD);
 REGISTER_TENSORRT_LAYER_BUILDER(BatchNorm, LAYER_SCALE);
 
 }  //  namespace TNN_NS