建议您添加如下代码
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)
因为这可能是tf自身的bug,官方issue中有写,但目前未处理的状态,评论中提供了一些解决方法。您可以尝试加一下以上代码,看能解决该问题。
https://github.com/tensorflow/tensorflow/issues/24828