bug fix: stack was not being unmasked before return in layer_op of DenseFeatureStackBlock
@@ -415,6 +415,23 @@ class DenseFeatureStackBlock(TrainableLayer):
stack.append(conv)
+ # unmask the stack
+ channels = tf.unstack(stack, axis=-1)
+ # exclude the input channels
+ channels = channels[input_channels:]
+ input_mask = input_mask[input_channels:]
+ # create a channel with zeros to be placed where channels were not calculated
+ channels.insert(0, tf.zeros_like(channels[0]))
+ # re stack
+ channels = tf.stack(channels, axis=-1)
+
+ # indices to keep
+ int_mask = tf.cast(input_mask, tf.int32)
+ indices = tf.cumsum(int_mask) * int_mask
+ # rearrange stack with zeros where channels were not calculated
+ stack = tf.gather(channels, indices, axis=-1)
+ stack = tf.unstack(stack, axis=-1)
+
return stack
