This problem remains in TensorFlow 2.10 and tensorflow-metal 0.6.0. I'm trying to train a 3D U-Net. The input to the batch norm layer is 5-dimensional, and I get the same error. The same model trains just fine on Windows.

I also get this issue when training a PointNet model (https://keras.io/examples/vision/pointnet/). My setup is tensorflow-macos 2.9.2, tensorflow-metal 0.5.0, Python 3.9.13 on a Mac Studio M1 Max with 64GB memory running macOS 12.4. For the timebeing I've reverted to an older Ubuntu machine for training. I could just use CPU training on the M1 I suppose. It would be great to get this resolved because I love using the M1 for everything. It's so fast and quiet.