'RaggedTensor' object has no attribute 'set_shape' when passing RaggedTensor to BatchNormalization layer in Tensorflow 2.8

Question

I am trying to pass a 3D RaggedTensor of shape (batch_size=None, 1, 136) to a batch normalization layer, but am receiving the error in the title. I assume because RaggedTensors can very in dimensions, and thus have no set_shape() method as it would be difficult to define how to coerce the sub tensors of various shapes to the target shape.

Is there a way to get this to work or do RaggedTensors just not work with batchnorm?

Thanks