The method refers to
https://github.com/tensorflow/tensorflow/issues/45070
In the main code add
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
to fix this error
quite simple