How to train on CPU in tensorflow 2.6 in mac M1?

Question

I recently update from Big Sur to Monterey, but my tensorflow was spoiled with error of memory allocation (malloc issues)

I reinstall Tensorflow following these instructions and it seems to work but I search the way to train on 'CPU' or in 'Neural Engine' but not found.

In "Big Sur" OS, my code was like

from tensorflow.python.compiler.mlcompute import mlcompute
mlcompute.set_mlc_device(device_name='any')

In my experiments I obtain better results if train on small batches and train on 'ANY' was a lot faster that train on 'GPU'

Is there something equivalent for train on 'CPU' or 'ANY' in tensorflow 2.6 on M1 mac and OS Monterey?

Thanks in advance!