Michael_Sidorov:

I do state the shapes of the images and the labels


Ah, yes that could use a little more explanation.
tf.data and tf.function (also used in keras) use TensorFlow’s 2-stage process. First it builds the graph, and then it executes the graph on your tensors.
Tensor.shape is the shape that was determined at graph build time (tf.shape(tensor) gets you the runtime shape).
So what’s happening here is that tf.numpy_function(load_image, ...)  says “don’t build a graph for this part, just run it in python”. So none of the code in  load_image runs during graph building, and TensorFlow doesn’t know you’re setting the shape in there.
Try putting only the non-tensorflow parts in your tf.numpy_function(load_image, ...). Move everything from img = preprocessing_func(image=img) down, into another function, then you should be able to ds.map() that function without the numpy_function wrapper.
But if you can use the tfio’s decode tiff, maybe you can drop the numpy_function entierly.