tensorflow2.4.1 无法使用gpu,只能用cpu

error:

E tensorflow/stream_executor/cuda/cuda_blas.cc:226] failed to create
cublas handle: CUBLAS_STATUS_ALLOC_FAILED
tensorflow.python.framework.errors_impl.InternalError: Blas GEMM launch failed : a.shape=(100, 2048)

跑训练的时候发现只能用cpu跑，gpu根本不动

后面发现是session写的有问题
原写法
with tf.compat.v1.Session() as sess:

改成

config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.compat.v1.Session(config=config)
with session as sess:

成功启动gpu

ps:我是用tensorflow2跑的tensorflow代码，不知道是否和这个有关

TensorFlow 1.X 的 Graph Execution 下，可以在实例化新的session时传入 tf.compat.v1.ConfigPhoto 类来设置TensorFlow使用显存的策略。具体方式是实例化一个 tf.ConfigProto 类，设置参数，并在创建 tf.compat.v1.Session 时指定Config参数。以下代码通过 allow_growth 选项设置TensorFlow仅在需要时申请显存空间