M1 Max 64GB, Monterey 12.0.1 This is a strange error. To resolve I added the expected input shape to the first layer # Failing: # Possibly unrelated: the `with` block was added as the following code didn't run on GPU: with tf.device('/cpu:0'): data_augmentation = tf.keras.Sequential([ tf.keras.layers.RandomFlip('horizontal'), tf.keras.layers.RandomRotation(0.2), ]) The above caused the following errors # The following error was received - but after adding the input shape it worked fine # 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_=[] # RngReadAndSkip: CPU # _Arg: GPU CPU # Colocation members, user-requested devices, and framework assigned devices, if any: # model_sequential_2_random_flip_2_stateful_uniform_full_int_rngreadandskip_resource (_Arg) framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0 # model/sequential_2/random_flip_2/stateful_uniform_full_int/RngReadAndSkip (RngReadAndSkip) # [[{{node model/sequential_2/random_flip_2/stateful_uniform_full_int/RngReadAndSkip}}]] [Op:__inference_train_function_12915] Adding the input_shape size seems to have resolved the issue and the model is training. # Working with tf.device('/cpu:0'): data_augmentation = tf.keras.Sequential([ tf.keras.layers.RandomFlip('horizontal',input_shape=(IMG_SIZE[0],IMG_SIZE[1],3)), tf.keras.layers.RandomRotation(0.2), ]) I also had this configuration: tf.config.set_soft_device_placement(True)