I found a workaround to make AdamW work on Apple Silicon using the latest version of tensorflow, tensorflow-addons. All you need to do is to import AdamW from tensorflow_addons.optimizers and you should be good. I'm not using it with tensorflow-metal though, there's a huge impact in performance. (At least 4x slower) 👎 (btw, the msg "There's a known slowdown on M1/M2 devices.... falling back.... " won't appear anymore) Hope this helps! 🙂

Hi there! I've been getting the same error while using RectifiedAdam and Lookahead on my m1 max. Traceback (most recent call last): File "/Users/netanel/dev/models/transformer_regressor.py", line 105, in fit model.fit(train_X, train_y, epochs=200, batch_size=128, File "/opt/homebrew/Caskroom/miniconda/base/envs/dl/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler raise e.with_traceback(filtered_tb) from None File "/opt/homebrew/Caskroom/miniconda/base/envs/dl/lib/python3.10/site-packages/tensorflow/python/eager/execute.py", line 53, in quick_execute tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device for operation model/multi_head_attention/query/einsum/Einsum/ReadVariableOp: Could not satisfy explicit device specification '' because the node {{colocation_node model/multi_head_attention/query/einsum/Einsum/ReadVariableOp}} was colocated with a group of nodes that required incompatible device '/job:localhost/replica:0/task:0/device:GPU:0'. All available devices [/job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:GPU:0]. Do we have any news?