ValueError: When using data tensors as input to a model

I'm trying to train a model with GradientTape in Keras. Here is the code:

@tf.function
def train_step(x,y):
    
    with tf.GradientTape() as tape:
                
        predictions = model.predict(x)
        
        loss = compute_loss(y, predections)
    
    grads = tape.gradient(loss, model.trainable_variables)
        
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    
    return loss

history = []

for iter in tqdm(range(num_iters)):
    
    x_batch, y_batch = get_batch(x_train, y_train, batch_dim)
    loss = train_step(x_batch, y_batch)
    history.append(loss.numpy().mean())
    

This code leads to the following error:

ValueError: When using data tensors as input to a model, you should specify the `steps` argument.

However if I try to call the prediction outside the function as follows:

history = []

for iter in tqdm(range(num_iters)):
    
    x_batch, y_batch = get_batch(x_train, y_train, batch_dim)       
    x_hat = model.predict(x_batch)

I get no error...

Can someone explain me why do I get this behavior from Keras?