tf.image.decode_jpeg often taking forever to load file

The following code is part of my code for a tf graph to read images. When I use this code to iterate through the data, the program gets stuck in tf.io.read_file(path) after a few hundred images forever and doesn't do anything. More specifically, the code even can't be paused and I had to restart the session every time.

@tf.function()
def read_image(path):
  image = tf.io.read_file(path)
  image = tf.image.decode_jpeg(image)
  return image

...

div8k_list=[os.path.join(div8k_save_path, x) for x in os.listdir(div8k_save_path)]
train_path = tf.data.Dataset.from_tensor_slices(div8k_list)

train_images = train_path.map(read_image, num_parallel_calls=tf.data.AUTOTUNE)

I first suspected that there were a few corrupted images or wrong paths in the data that were causing this problem and tested the following code.

for path in train_path:
  print(path)
  
  image = tf.io.read_file(path)
  image = tf.image.decode_jpeg(image)

Surprisingly, there was no common characteristic of the image path the loop was stuck. And it was not a problem of the image because the loop was once stuck at 1056.png but when I explicitly loaded 1056.png, there was no problem.

What could be the cause of this problem?

edit: to summarize, the program is stuck at read_image forever, while I couldn't find a problem in the dataset.

My dataset is the DIV8K dataset and I am running in COLAB.

EDIT The function that is slowing my code is decode_jpeg, because the following definition of read_image worked multiple times.

@tf.function()
def read_image(path):
  image = tf.io.read_file(path)
  image = tf.image.decode_jpeg(image)
  return image