Saving Custom Model Not Working
@aecelaya,
If your model contains custom objects or custom layers then you need to define **get_config** and **from_config** methods in the Model. Then during model.save() you need to pass an argument to custom_objects with dictionary of custom objects.

Please refer to attached [source](https://www.tensorflow.org/guide/keras/save_and_serialize#defining_the_config_methods) for a minimal example.

You may also refer to the keras https://github.com/keras-team/tf-keras/issues/64. Please try to recreate the model accordingly.

Thanks! Thank you for the reply! I tried what you suggested, but I still see the same error.

For now, I can move forward without needing to save the model. I'll revisit this issue later.

Thank you! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60022">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60022">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Support

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

tf 2.11.0

### Custom Code

Yes

### OS Platform and Distribution

Colab Notebook

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

```shell
I'm trying to save a custom model where the number of channels increases as we progress along the forward pass of the architecture. The forward pass of the model works as expected, and I can train it. However, when trying to save the model, I get an error concerning the number of channels passed to certain convolutional layers. 

Is there a way around this issue?
```


### Standalone code to reproduce the issue

```shell
https://colab.research.google.com/drive/1YZd8OQWjgYasVblDit-gRNvA-kFJdRkn?usp=sharing
```


### Relevant log output

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-7-696ba1b852d4> in <module>
----> 1 wnet.save("/content/wnet")

7 frames
/tmp/__autograph_generated_filelz68b5h5.py in tf__call(self, x)
      8                 do_return = False
      9                 retval_ = ag__.UndefinedReturnValue()
---> 10                 x = ag__.converted_call(ag__.ld(self).conv, (ag__.ld(x),), None, fscope)
     11                 x = ag__.converted_call(ag__.ld(self).norm, (ag__.ld(x),), None, fscope)
     12                 x = ag__.converted_call(ag__.ld(self).activation, (ag__.ld(x),), None, fscope)

ValueError: Exception encountered when calling layer 'base_model' (type BaseModel).

in user code:

    File "<ipython-input-1-355bdcd220c4>", line 149, in call  *
        x = decoder_block([*previous_skips[str(self.global_depth - 1 - i)],
    File "/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_file30lnbd33.py", line 12, in tf__call
        x = ag__.converted_call(ag__.ld(self).block, (ag__.ld(x),), None, fscope)
    File "/tmp/__autograph_generated_filejlndlto_.py", line 10, in tf__call
        x = ag__.converted_call(ag__.ld(self).conv1, (ag__.ld(x),), None, fscope)
    File "/tmp/__autograph_generated_filelz68b5h5.py", line 10, in tf__call
        x = ag__.converted_call(ag__.ld(self).conv, (ag__.ld(x),), None, fscope)

    ValueError: Exception encountered when calling layer 'decoder_block_1' (type DecoderBlock).
    
    in user code:
    
        File "<ipython-input-1-355bdcd220c4>", line 96, in call  *
            x = self.block(x)
        File "/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
            raise e.with_traceback(filtered_tb) from None
        File "/tmp/__autograph_generated_filejlndlto_.py", line 10, in tf__call
            x = ag__.converted_call(ag__.ld(self).conv1, (ag__.ld(x),), None, fscope)
        File "/tmp/__autograph_generated_filelz68b5h5.py", line 10, in tf__call
            x = ag__.converted_call(ag__.ld(self).conv, (ag__.ld(x),), None, fscope)
    
        ValueError: Exception encountered when calling layer 'u_net_block_7' (type UNetBlock).
        
        in user code:
        
            File "<ipython-input-2-76991f1be87b>", line 14, in call  *
                x = self.conv1(x)
            File "/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
                raise e.with_traceback(filtered_tb) from None
            File "/tmp/__autograph_generated_filelz68b5h5.py", line 10, in tf__call
                x = ag__.converted_call(ag__.ld(self).conv, (ag__.ld(x),), None, fscope)
        
            ValueError: Exception encountered when calling layer 'conv_layer_14' (type ConvLayer).
            
            in user code:
            
                File "<ipython-input-1-355bdcd220c4>", line 56, in call  *
                    x = self.conv(x)
                File "/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
                    raise e.with_traceback(filtered_tb) from None
                File "/usr/local/lib/python3.9/dist-packages/keras/engine/input_spec.py", line 277, in assert_input_compatibility
                    raise ValueError(
            
                ValueError: Input 0 of layer "conv3d_14" is incompatible with the layer: expected axis -1 of input shape to have value 24, but received input with shape (None, 16, 16, 16, 32)
            
            
            Call arguments received by layer 'conv_layer_14' (type ConvLayer):
              • x=tf.Tensor(shape=(None, 16, 16, 16, 32), dtype=float32)
        
        
        Call arguments received by layer 'u_net_block_7' (type UNetBlock):
          • x=tf.Tensor(shape=(None, 16, 16, 16, 32), dtype=float32)
    
    
    Call arguments received by layer 'decoder_block_1' (type DecoderBlock):
      • skip=['tf.Tensor(shape=(None, 16, 16, 16, 8), dtype=float32)', 'tf.Tensor(shape=(None, 16, 16, 16, 8), dtype=float32)', 'tf.Tensor(shape=(None, 16, 16, 16, 8), dtype=float32)']
      • x=tf.Tensor(shape=(None, 8, 8, 8, 8), dtype=float32)


Call arguments received by layer 'base_model' (type BaseModel):
  • x=tf.Tensor(shape=(None, 8, 8, 8, 8), dtype=float32)
  • previous_skips={'0': ListWrapper(['tf.Tensor(shape=(None, 64, 64, 64, 8), dtype=float32)']), '1': ListWrapper(['tf.Tensor(shape=(None, 32, 32, 32, 8), dtype=float32)']), '2': ListWrapper(['tf.Tensor(shape=(None, 16, 16, 16, 8), dtype=float32)', 'tf.Tensor(shape=(None, 16, 16, 16, 8), dtype=float32)'])}
  • previous_peaks={'0': ListWrapper([]), '1': ListWrapper(['tf.Tensor(shape=(None, 32, 32, 32, 8), dtype=float32)']), '2': ListWrapper(['tf.Tensor(shape=(None, 16, 16, 16, 8), dtype=float32)'])}
```
</details>