Retinanet build issue with Jax

[divyashreepathihalli]

error dump :

/tmpfs/venv/lib/python3.9/site-packages/keras/src/utils/traceback_utils.py:122: in error_handler
    raise e.with_traceback(filtered_tb) from None
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = 
iterator = 
data_batch = {'bounding_boxes': {'boxes': , 'cla...s_tensor_29218>}, 'images': }

    def _symbolic_build(self, iterator=None, data_batch=None):
        model_unbuilt = not all(layer.built for layer in self._flatten_layers())
        compile_metrics_unbuilt = (
            self._compile_metrics is not None
            and not self._compile_metrics.built
        )
        compile_loss_unbuilt = (
            self._compile_loss is not None and not self._compile_loss.built
        )
        optimizer_unbuilt = (
            self.optimizer is not None and not self.optimizer.built
        )
        if model_unbuilt or compile_metrics_unbuilt or compile_loss_unbuilt:
            # Create symbolic tensors matching an input batch.
    
            def to_symbolic_input(v):
                if v is None:
                    return None
                return backend.KerasTensor(
                    v.shape, backend.standardize_dtype(v.dtype)
                )
    
            if data_batch is None:
                for _, data in iterator.enumerate_epoch():
                    data_batch = data[0]
                    break
            data_batch = tree.map_structure(to_symbolic_input, data_batch)
            (
                x,
                y,
                sample_weight,
            ) = data_adapter_utils.unpack_x_y_sample_weight(data_batch)
    
            # Build all model state with `backend.compute_output_spec`.
            try:
                y_pred = backend.compute_output_spec(self, x, training=False)
            except Exception as e:
>               raise RuntimeError(
                    "Unable to automatically build the model. "
                    "Please build it yourself before calling "
                    "fit/evaluate/predict. "
                    "A model is 'built' when its variables have "
                    "been created and its `self.built` attribute "
                    "is True. Usually, calling the model on a batch "
                    "of data is the right way to build it.\n"
                    "Exception encountered:\n"
                    f"'{e}'"
                )
E               RuntimeError: Unable to automatically build the model. Please build it yourself before calling fit/evaluate/predict. A model is 'built' when its variables have been created and its `self.built` attribute is True. Usually, calling the model on a batch of data is the right way to build it.
E               Exception encountered:
E               'Exception encountered when calling RetinaNet.call().
E               
E               Invalid input shape for input Tracedwith. Expected shape (None, None, None, 3), but input has incompatible shape (5, 3, 4)
E               
E               Arguments received by RetinaNet.call():
E                 • inputs={'bounding_boxes': {'boxes': 'jnp.ndarray(shape=(5, 3, 4), dtype=float32)', 'classes': 'jnp.ndarray(shape=(5, 3), dtype=float32)', 'num_dets': 'jnp.ndarray(shape=(5,), dtype=float32)'}, 'images': 'jnp.ndarray(shape=(5, 512, 512, 3), dtype=float32)'}
E                 • training=False
E                 • mask={'bounding_boxes': {'boxes': 'None', 'classes': 'None', 'num_dets': 'None'}, 'images': 'None'}