src/transformers/models/hubert/modeling_tf_hubert.py/TFHubertGroupNorm/_create_input_spec
class TFHubertGroupNorm:    def _create_input_spec(self, input_shape):
        dim = input_shape[self.axis]
        self.input_spec = tf.keras.layers.InputSpec(ndim=len(input_shape), axes={self.axis: dim})