Ranking model load weights after checkpointing failing

Hi ,

I am training the ranking model and using tensorflow callbacks to save the checkpoints .
I did it but when loading the same model it loads but while predicting it throws error as below -

InvalidArgumentError: In[0] mismatch In[1] shape: 32 vs. 64: [32764,32] [64,256] 0 0
[[node f2m_ranking/ranking_model/sequential_2/dense/Tensordot/MatMul (defined at tmp/ipykernel_21470/657321342.py:39) ]] [Op:__inference_predict_function_630]

InvalidArgumentError Traceback (most recent call last)
/tmp/ipykernel_21470/4271407944.py in
----> 1 model.predict(test_full,verbose=1)

/opt/conda/lib/python3.7/site-packages/keras/engine/training.py in predict(self, x, batch_size, verbose, steps, callbacks, max_queue_size, workers, use_multiprocessing)
1749 for step in data_handler.steps():
1750 callbacks.on_predict_batch_begin(step)
-> 1751 tmp_batch_outputs = self.predict_function(iterator)
1752 if data_handler.should_sync:
1753 context.async_wait()

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in call(self, *args, **kwds)
883
884 with OptionalXlaContext(self._jit_compile):
--> 885 result = self._call(*args, **kwds)
886
887 new_tracing_count = self.experimental_get_tracing_count()

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
948 # Lifting succeeded, so variables are initialized and we can run the
949 # stateless function.
--> 950 return self._stateless_fn(*args, **kwds)
951 else:
952 _, _, _, filtered_flat_args = \

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in call(self, *args, **kwargs)
3038 filtered_flat_args) = self._maybe_define_function(args, kwargs)
3039 return graph_function._call_flat(
-> 3040 filtered_flat_args, captured_inputs=graph_function.captured_inputs) # pylint: disable=protected-access
3041
3042 @property

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _call_flat(self, args, captured_inputs, cancellation_manager)
1962 # No tape is watching; skip to running the function.
1963 return self._build_call_outputs(self._inference_function.call(
-> 1964 ctx, args, cancellation_manager=cancellation_manager))
1965 forward_backward = self._select_forward_and_backward_functions(
1966 args,

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in call(self, ctx, args, cancellation_manager)
594 inputs=args,
595 attrs=attrs,
--> 596 ctx=ctx)
597 else:
598 outputs = execute.execute_with_cancellation(

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
58 ctx.ensure_initialized()
59 tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 60 inputs, attrs, num_outputs)
61 except core._NotOkStatusException as e:
62 if name is not None:

**_**InvalidArgumentError: In[0] mismatch In[1] shape: 32 vs. 64: [32764,32] [64,256] 0 0
[[node f2m_ranking/ranking_model/sequential_2/dense/Tensordot/MatMul (defined at tmp/ipykernel_21470/657321342.py:39) ]] [Op:__inference_predict_function_630]

Function call stack:
predict_function***_*
​