“Cannot add tensor to the batch: number of elements does not match” while iterating through dataset

I included the code (up to the iterating part) in a Colab notebook here. I fed my data to my train_dataset like this.

batch_size = 30

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = (
    train_dataset.map(encode_single_sample, num_parallel_calls=tf.data.experimental.AUTOTUNE)
                 .batch(batch_size) #this one gives the error
                 .prefetch(buffer_size=tf.data.experimental.AUTOTUNE))

I get a InvalidArgumentError: Cannot add tensor to the batch: number of elements does not match. Shapes are: [tensor]: [78], [batch]: [34] when I iterate through the dataset:

for batch in train_dataset.take(1): #error while iterating

The .batch(batch_size) above gives the error. I noticed that the [tensor] and [batch] numbers in the error message change every time I run the code again. The code runs when I change batch_size to 1, but I'd like to be able to change it as well.

Would appreciate it a lot if anyone could debug this part

Edit 1: I'm following Keras' OCR tutorial here, and their code runs just fine.

from “Cannot add tensor to the batch: number of elements does not match” while iterating through dataset