Device assignment issue

I've installed the tensorflow-macos and tensorflow-metal on my M1 Mac mini and I wanted to try training a LSTM network, but I got this error shown below. I think it's related to the assignment of GPU and CPU. The macOS version is 11.6. I will greatly appreciate any help. Thank you!

InvalidArgumentError: Cannot assign a device for operation sequential_4/lstm_16/Read/ReadVariableOp: Could not satisfy explicit device specification '' because the node {{colocation_node sequential_4/lstm_16/Read/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' assigned_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' resource_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' supported_device_types_=[CPU] possible_devices_=[]
ResourceApplyAdaMax: CPU
ReadVariableOp: GPU CPU
_Arg: GPU CPU
Colocation members, user-requested devices, and framework assigned devices, if any:
sequential_4_lstm_16_read_readvariableop_resource (_Arg) framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0
adamax_adamax_update_resourceapplyadamax_m (_Arg) framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0
adamax_adamax_update_resourceapplyadamax_v (_Arg) framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0
sequential_4/lstm_16/Read/ReadVariableOp (ReadVariableOp)
Adamax/Adamax/update/ResourceApplyAdaMax (ResourceApplyAdaMax) /job:localhost/replica:0/task:0/device:GPU:0
[[{{node sequential_4/lstm_16/Read/ReadVariableOp}}]] [Op:__inference_train_function_75386]

I'm also having similar problems, but Apple just does not give any support!

Hi @chuhaodeng,

Thanks for reporting this. It has to do with the ResourceApplyAdaMax op missing from the Metal plugin based on the error message. We will be taking a look at this issue and update here when there are any news on this.

Device assignment issue
 
 
Q