title: Saving model on Tensorflow 2.7.0 with data augmentation layer
body: <p>I am getting an error when trying to save a model with data augmentation layers with Tensorflow version 2.7.0.</p>
  <p>Here is the code of data augmentation:</p>
  <pre><code>input_shape_rgb = (img_height, img_width, 3)
  data_augmentation_rgb = tf.keras.Sequential(
    [ 
      layers.RandomFlip(&quot;horizontal&quot;),
      layers.RandomFlip(&quot;vertical&quot;),
      layers.RandomRotation(0.5),
      layers.RandomZoom(0.5),
      layers.RandomContrast(0.5),
      RandomColorDistortion(name='random_contrast_brightness/none'),
    ]
  )
  </code></pre>
  <p>Now I build my model like this:</p>
  <pre><code># Build the model
  input_shape = (img_height, img_width, 3)
  
  model = Sequential([
    layers.Input(input_shape),
    data_augmentation_rgb,
    layers.Rescaling((1./255)),
  
    layers.Conv2D(16, kernel_size, padding=padding, activation='relu', strides=1, 
       data_format='channels_last'),
    layers.MaxPooling2D(),
    layers.BatchNormalization(),
  
    layers.Conv2D(32, kernel_size, padding=padding, activation='relu'), # best 4
    layers.MaxPooling2D(),
    layers.BatchNormalization(),
  
    layers.Conv2D(64, kernel_size, padding=padding, activation='relu'), # best 3
    layers.MaxPooling2D(),
    layers.BatchNormalization(),
  
    layers.Conv2D(128, kernel_size, padding=padding, activation='relu'), # best 3
    layers.MaxPooling2D(),
    layers.BatchNormalization(),
  
    layers.Flatten(),
    layers.Dense(128, activation='relu'), # best 1
    layers.Dropout(0.1),
    layers.Dense(128, activation='relu'), # best 1
    layers.Dropout(0.1),
    layers.Dense(64, activation='relu'), # best 1
    layers.Dropout(0.1),
    layers.Dense(num_classes, activation = 'softmax')
   ])
  
   model.compile(loss='categorical_crossentropy', optimizer='adam',metrics=metrics)
   model.summary()
  </code></pre>
  <p>Then after the training is done I just make:</p>
  <pre><code>model.save(&quot;./&quot;)
  </code></pre>
  <p>And I'm getting this error:</p>
  <pre><code>---------------------------------------------------------------------------
  KeyError                                  Traceback (most recent call last)
  &lt;ipython-input-84-87d3f09f8bee&gt; in &lt;module&gt;()
  ----&gt; 1 model.save(&quot;./&quot;)
  
  
  /usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py in 
   error_handler(*args, **kwargs)
   65     except Exception as e:  # pylint: disable=broad-except
   66       filtered_tb = _process_traceback_frames(e.__traceback__)
   ---&gt; 67       raise e.with_traceback(filtered_tb) from None
   68     finally:
   69       del filtered_tb
  
   /usr/local/lib/python3.7/dist- 
   packages/tensorflow/python/saved_model/function_serialization.py in 
   serialize_concrete_function(concrete_function, node_ids, coder)
   66   except KeyError:
   67     raise KeyError(
   ---&gt; 68         f&quot;Failed to add concrete function '{concrete_function.name}' to 
   object-&quot;
   69         f&quot;based SavedModel as it captures tensor {capture!r} which is 
   unsupported&quot;
   70         &quot; or not reachable from root. &quot;
  
   KeyError: &quot;Failed to add concrete function 
   'b'__inference_sequential_46_layer_call_fn_662953'' to object-based SavedModel as it 
   captures tensor &lt;tf.Tensor: shape=(), dtype=resource, value=&lt;Resource Tensor&gt;&gt; which 
   is unsupported or not reachable from root. One reason could be that a stateful 
   object or a variable that the function depends on is not assigned to an attribute of 
   the serialized trackable object (see SaveTest.test_captures_unreachable_variable).&quot;
  </code></pre>
  <p>I inspected the reason of getting this error by changing the architecture of my model and I just found that reason came from the data_augmentation layer since the <code>RandomFlip</code> and <code>RandomRotation</code> and others are changed from <code>layers.experimental.prepocessing.RandomFlip</code> to <code>layers.RandomFlip</code>, but still the error appears.</p>