Make TensorFlow use the GPU on an ARM Mac

I have installed TensorFlow on an M1 (ARM) Mac according to these instructions. Everything works fine.

However, model training is happening on the CPU. How do I switch training to the GPU?

In: tensorflow.config.list_physical_devices()
Out: [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

In the documentation of Apple's TensorFlow distribution I found the following slightly confusing paragraph:

It is not necessary to make any changes to your existing TensorFlow scripts to use ML Compute as a backend for TensorFlow and TensorFlow Addons. There is an optional mlcompute.set_mlc_device(device_name='any') API for ML Compute device selection. The default value for device_name is 'any', which means ML Compute will select the best available device on your system, including multiple GPUs on multi-GPU configurations. Other available options are CPU and GPU. Please note that in eager mode, ML Compute will use the CPU. For example, to choose the CPU device, you may do the following:

# Import mlcompute module to use the optional set_mlc_device API for device selection with ML Compute.
from tensorflow.python.compiler.mlcompute import mlcompute

# Select CPU device.
mlcompute.set_mlc_device(device_name='cpu') # Available options are 'cpu', 'gpu', and 'any'.

So I try to run:

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

and get:

WARNING:tensorflow: Eager mode uses the CPU. Switching to the CPU.

At this point I am stuck. How can I train keras models on the GPU to my MacBook Air?

TensorFlow version: 2.4.0-rc0

from Make TensorFlow use the GPU on an ARM Mac