Update assertions in batch norm layer

artoolkitx · Mar 12, 2018 · ab20d2a · ab20d2a
1 parent 69a8f11
commit ab20d2a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/modules/dnn/src/layers/batch_norm_layer.cpp b/modules/dnn/src/layers/batch_norm_layer.cpp
@@ -32,7 +32,7 @@ class BatchNormLayerImpl : public BatchNormLayer
     BatchNormLayerImpl(const LayerParams& params)
     {
         setParamsFrom(params);
-        CV_Assert(blobs.size() >= 3);
+        CV_Assert(blobs.size() >= 2);
 
         hasWeights = params.get<bool>("has_weight", false);
         hasBias = params.get<bool>("has_bias", false);
@@ -46,8 +46,8 @@ class BatchNormLayerImpl : public BatchNormLayer
                   blobs[0].type() == CV_32F && blobs[1].type() == CV_32F);
 
         float varMeanScale = 1.f;
-        if (!hasWeights && !hasBias) {
-            CV_Assert(blobs[2].type() == CV_32F);
+        if (!hasWeights && !hasBias && blobs.size() > 2) {
+            CV_Assert(blobs.size() == 3, blobs[2].type() == CV_32F);
             varMeanScale = blobs[2].at<float>(0);
             if (varMeanScale != 0)
                 varMeanScale = 1/varMeanScale;