Saved and loaded subclassed keras model throws ValueError "Could not find matching function" for non-default call args tensorflow 2

I am trying to load a custom sub-classed keras model and to then call it on a new batch.

However, the call on the loaded model throws the error "ValueError: Could not find matching function to call loaded from the SavedModel." when it is called using args that are different in value than the defaults. When the call uses the default arg values, the model does the inference without an issue. Further, before the save and reload, the model can do the call with both default and non-default arg values. One final note, the arg "training" seems to work fine - it is only when I add a new arg ("full_batch_eval" in the example below) that this issue happens.

Below is a self contained example demonstrating the issue. Any idea about what this issue is? I'm guessing that the call function isn't traced during the save for all values of the args and thus these aren't available in the loaded version of the model? Thanks in advance!

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Flatten, Dense

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()
x_train = tf.expand_dims(x_train.astype("float32")/255.,-1)
x_test = tf.expand_dims(x_test.astype("float32")/255.,-1)
y_train = y_train.astype("float32")
y_test = y_test.astype("float32")

batch_size = 32
buffer_size = 1024
train_ds = tf.data.Dataset.from_tensor_slices((x_train,y_train))
train_ds = train_ds.shuffle(buffer_size=buffer_size).batch(batch_size)

class TestModel(keras.Model):
    def __init__(self,**kwargs):
        super(TestModel,self).__init__(**kwargs)
        self.f = Flatten()
        self.dense_1 = Dense(units=10,activation=tf.nn.relu)
    
    def call(self,inputs,training=False,full_batch_eval=True):
        x = inputs.get('x')
        return self.dense_1(self.f(x))

def trainTestModel(m,ds,l,opt,epochs=1):
    @tf.function
    def train_step(x,y):
        with tf.GradientTape() as t0:
            out0 = m({'x':x,'y':y})
            l0 = l(y_true=y,y_pred=out0)
            g0 = t0.gradient(target=l0,sources=m.trainable_variables)
            opt.apply_gradients(zip(g0,m.trainable_variables))
    
    for epoch in range(epochs):
        for (x,y) in ds:
            train_step(x,y)

m = TestModel()

l = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
opt = tf.keras.optimizers.Adam(learning_rate=1e-4)
trainTestModel(m=m,ds=train_ds.take(1000),l=l,opt=opt)

for (x,y) in train_ds.take(1):
    t = m({'x':x,'y':y}) #works
    t = m({'x':x,'y':y},training=True) #works
    t = m({'x':x,'y':y},training=False,full_batch_eval=True) #works
    t = m(inputs={'x':x,'y':y},training=True,full_batch_eval=False) #works
    t = m(inputs={'x':x,'y':y},training=True,full_batch_eval=True) #works
    t = m(inputs={'x':x,'y':y},training=False,full_batch_eval=False) #works

prefix = 'saved_model_dir'
tf.saved_model.save(m,export_dir=prefix)

m = tf.saved_model.load(prefix)

for (x,y) in train_ds.take(1):
    t = m({'x':x,'y':y}) #works
    t = m({'x':x,'y':y},training=True) #works
    t = m({'x':x,'y':y},training=False,full_batch_eval=True) #works
    t = m(inputs={'x':x,'y':y},training=True,full_batch_eval=False) #error
    t = m(inputs={'x':x,'y':y},training=True,full_batch_eval=True) #works
    t = m(inputs={'x':x,'y':y},training=False,full_batch_eval=False) #error