How to fix 'Mul' Op has type float32 that does not match type float64?

Question

diffusion_ft_trainer.fit(training_dataset, epochs=epochs, callbacks=[ckpt_callback])

TypeError: Input 'y' of 'Mul' Op has type float32 that does not match type float64 of argument 'x'.

training_dataset
<_PrefetchDataset element_spec={'images': TensorSpec(shape=(None, 256, 256, 3), dtype=tf.float32, name=None), 'tokens': TensorSpec(shape=(None, 77), dtype=tf.float64, name=None), 'encoded_text': TensorSpec(shape=(None, 77, 768), dtype=tf.float32, name=None)}

i tried to convert image into tf.float64 using tf.cast method

training_dataset = training_dataset.map(lambda x: {'images': tf.cast(x['images'], tf.float64), 'tokens': x['tokens'], 'encoded_text': x['encoded_text']})

but i am getting same error even if i converted images to type 64

here is link to my notebook text