Dataset map function error : TypeError: Expected list for 'input' argument to 'EagerPyFunc' Op, not Tensor

I am currently trying to write a script to create a TFRecord file.

Therefore, I am following the instruction on the offical tensorflow website: https://www.tensorflow.org/tutorials/load_data/tfrecord#writing_a_tfrecord_file

However, when applying the map function to each element of the Dataset I get an error that I do not understand.

This is my code (should be copy and pasteable):

import numpy as np
import tensorflow as tf
from tensorflow.data import Dataset


def generate_random_img_data(n_count=10, patch_size=5):
    return np.random.randint(low=0, high=256, size=(n_count, patch_size, patch_size, 3))


def as_int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def serialize_one_image(img):
    features = {}
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            for k in range(img.shape[2]):
                features.update({str(i) + _ + str(j) + _ + str(k) : as_int64_feature(img[i,j,k]) })
    example_proto = tf.train.Example(features=tf.train.Features(feature=features))
    return example_proto.SerializeToString()


def tf_serialize_one_image(img):
    tf_string = tf.py_function(serialize_one_image, img, tf.string)
    return tf.reshape(tf_string,())


ds = Dataset.from_tensor_slices(generate_random_img_data())
ds_serialized = ds.map(tf_serialize_one_image) # --- not working

And when running this code I get the error:

TypeError: in user code:

    ipython-input-116-ec81a7077c70:25 tf_serialize_one_image  *
        tf_string = tf.py_function(serialize_one_image, img, tf.string)
    /Users/Tom/ML-Projects/vdst/lib/python3.7/site-packages/tensorflow/python/ops/script_ops.py:455 eager_py_func  **
        func=func, inp=inp, Tout=Tout, eager=True, name=name)
    /Users/Tom/ML-Projects/vdst/lib/python3.7/site-packages/tensorflow/python/ops/script_ops.py:341 _internal_py_func
        name=name)
    /Users/Tom/ML-Projects/vdst/lib/python3.7/site-packages/tensorflow/python/ops/gen_script_ops.py:69 eager_py_func
        name=name)
    /Users/Tom/ML-Projects/vdst/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py:385 _apply_op_helper
        (input_name, op_type_name, values))

    TypeError: Expected list for 'input' argument to 'EagerPyFunc' Op, not Tensor(args_0:0, shape=(5, 5, 3), dtype=int64).

What exactly am I doing wrong here?