I managed to solve this by installing Miniforge3 from https://github.com/conda-forge/miniforge for arm64 and following Apple instructions at https://developer.apple.com/metal/tensorflow-plugin/:
# Install Miniforge3
chmod +x ~/Downloads/Miniforge3-MacOSX-arm64.sh
bash Miniforge3-MacOSX-arm64.sh
conda --version
# Create venv
conda create -n ml-env python=3.9
conda activate ml-env
# Install TensorFlow
conda install -c apple tensorflow-deps
python -m pip install tensorflow-macos
python -m pip install tensorflow-metal
# Test the installation [https://developer.apple.com/metal/tensorflow-plugin/]
python test_tf.py
test_tf.py:
import tensorflow as tf
cifar = tf.keras.datasets.cifar100
(x_train, y_train), (x_test, y_test) = cifar.load_data()
model = tf.keras.applications.ResNet50(
include_top=True,
weights=None,
input_shape=(32, 32, 3),
classes=100,)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer="adam", loss=loss_fn, metrics=["accuracy"])
model.fit(x_train, y_train, epochs=5, batch_size=64)