tensorflow.math.reduce_std causes CUDA error during prediction but works fine during training

Question

I have a generator that loads some input data and a matching label. During training I call it through tf.data.Dataset.from_generator and map a preprocessing function with datase