GPU/CPU assignment error (just another thread)

I have a simple TCN model that I've been using for a while. Since my change to Apple M1, I am unable to run it.

My issue seems very similar to this and this, and I've also reported on FB9722799. On one of these threads Apple recognized that "we are aware of this issue and already working on a fix.". But this was 3 months ago, which is too much time without being able to fully use the computer for development!

If I uninstall tensorflow-metal, I can run the code (of course, only in CPU). If I install tensorflow-metal, I get the following error:

Cannot assign a device for operation model/conv_1_convolution/Conv2D/ReadVariableOp: Could not satisfy explicit device specification '' because the node {{colocation_node model/conv_1_convolution/Conv2D/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]. 
Colocation Debug Info:
Colocation group had the following types and supported devices: 
Root Member(assigned_device_name_index_=2 requested_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' 
...

	 [[{{node model/conv_1_convolution/Conv2D/ReadVariableOp}}]] [Op:__inference_train_function_15035]

Replies

Hi mrt77, can you please provide a code snippet to reproduce the issue?