import tensorflow as tf

            training=None,
            mask=None
    ):
        feature_hidden = inputs
        original_feature_hidden = inputs

        # flatten inputs
        if len(original_feature_hidden.shape) > 2:
            feature_hidden = tf.reshape(
                feature_hidden,
                [-1, feature_hidden.shape[-1]]
            )

        # pass it through fc_stack
        feature_hidden = self.fc_stack(
            feature_hidden,
