Bug fix ObjectAttentionBlock2D

Author: JanMarcelKezmann

tensorflow_advanced_segmentation_models/__version__.py

@@ -1,3 +1,3 @@
-VERSION = (0, 4, 6)
+VERSION = (0, 4, 7)
 
 __version__ = ".".join(map(str, VERSION))

tensorflow_advanced_segmentation_models/models/_custom_layers_and_blocks.py

@@ -883,7 +883,7 @@ def call(self, feats, ctx, training=None):
             ctx = tf.keras.layers.Permute((2, 1))(ctx)
             BS, H, W, C = feats.shape
             if self.scale > 1:
-                x = self.pool(x, training=training)
+                feats = self.pool(feats, training=training)
 
             query = self.f_pixel(feats, training=training)              # (BS, H, W, C)
             query = tf.keras.layers.Reshape((-1, C))(query)             # (BS, N, C)
@@ -906,7 +906,7 @@ def call(self, feats, ctx, training=None):
             # feats-dim: (BS, C, H, W) & ctx-dim: (BS, C, C2)
             BS, C, H, W = feats.shape
             if self.scale > 1:
-                x = self.pool(x, training=training)
+                feats = self.pool(feats, training=training)
 
             query = self.f_pixel(feats, training=training)              # (BS, C, H, W)
             query = tf.keras.layers.Reshape((C, -1))(query)             # (BS, C, N)