Tensorflow2.0 GPUconfig的问题

今天在把以前NMT代码迁到tf2.0 stable 上遇到了一个"CUDNN_STATUS_INTERNAL_ERROR"的问题,简单记录一下。
是一个跟gpu显存有关的问题,指定gpu动态分配即可。

Tensorflow2.0中gpuconfig相关的搬到了tf.config.experimental下
document

physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)

放在网络initialize前即可

github相关issue

你可能感兴趣的:(tensorflow)