TypeError: TextZeroShotDetectionModule.box_predictor() takes 1 positional argument but 3 were given

[Aki1991]

Hi all,
I am trying to convert my trained jax checkpoint in to tfmodel using this notebook provided at owl_vit/notebook directory.

While trying to convert the checkpoint to model, at this stage of the notebook,

assert tf.executing_eagerly(
), 'Use eager execution to avoid FailedPreconditionError.'
convert_and_save_model(
    predict_fn,
    variables,
    input_signatures=input_signatures,
    polymorphic_shapes=polymorphic_shapes,
    model_dir=EXPORT_DIR,
    with_gradient=True,
    enable_xla=True,
    compile_model=True)

I get this error.
TypeError: TextZeroShotDetectionModule.box_predictor() takes 1 positional argument but 3 were given

This happens when I use the official checkpoint (mentioned in the config file.) and also for the local checkpoint which I got after finetuning the model.