TensorFlow2.4.1 CUBLAS_STATUS_ALLOC_FAILED 問題解決方法

 在網路上找到的解決方法 .

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')

if gpus:

    try:

        # Currently, memory growth needs to be the same across GPUs

        for gpu in gpus:

            tf.config.experimental.set_memory_growth(gpu, True)

        logical_gpus = tf.config.experimental.list_logical_devices('GPU')

        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")

    except RuntimeError as e:

        # Memory growth must be set before GPUs have been initialized

        print(e)