fix avgpool imanagezono

intelligent-control-lab · Aug 4, 2024 · dd61366 · dd61366
1 parent d8e543f
commit dd61366
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/src/propagate/operators/pooling.jl b/src/propagate/operators/pooling.jl
@@ -42,9 +42,15 @@ also of type `ImageZonoBound`.
     type.
 """
 function propagate_layer(prop_method, layer::MeanPool, bound::ImageZonoBound, batch_info)
-    new_center = layer(bound.center)
-    new_generators = layer(bound.generators)
-    return ImageZonoBound(new_center, new_generators)
+    if size(bound.generators,4) > 0
+        new_center = layer(bound.center)
+        new_generators = layer(bound.generators)
+        return ImageZonoBound(new_center, new_generators)
+    else
+        new_center = layer(bound.center)
+        new_generators = zero(new_center)[:,:,:,1:0]
+        return ImageZonoBound(layer(bound.center), new_generators)
+    end
 end
 
 function propagate_layer_batch(prop_method::IBP, layer::MeanPool, bound::IBPBound, batch_info)