GPUを使おうとするとエラーが発生する

生じている問題

tensorflowでGPUを使おうとするとエラーが発生する

動作環境

・Python　3.9.11
・TensorFlow　2.8.0
・CUDA　11.4.4
・cuDNN　8.2.4
・GPU　RTX A6000
・NVIDIAドライバー　514.08
・VRAM　48571 MB

エラーメッセージ

学習を走らせようとすると、下のエラーが発生します。

2023-10-05 16:02:07.768555: I tensorflow/stream_executor/cuda/cuda_dnn.cc:368] Loaded cuDNN version 8204 2023-10-05 16:02:22.315203: E tensorflow/stream_executor/cuda/cuda_driver.cc:1163] failed to enqueue async memcpy from device to host: CUDA_ERROR_LAUNCH_TIMEOUT: the launch timed out and was terminated; host dst: 0x3d7d3db5c0; GPU src: 0x2f60850800; size: 8=0x8 2023-10-05 16:02:22.316034: I tensorflow/stream_executor/stream.cc:4442] [stream=000002328328CF00,impl=000002329929D2E0] INTERNAL: stream did not block host until done; was already in an error state 2023-10-05 16:02:22.316273: E tensorflow/stream_executor/cuda/cuda_driver.cc:756] failed to free device memory at 0000002F60850800; result: CUDA_ERROR_LAUNCH_TIMEOUT: the launch timed out and was terminated 2023-10-05 16:02:22.317091: W tensorflow/core/kernels/gpu_utils.cc:70] Failed to check cudnn convolutions for out-of-bounds reads and writes with an error message: 'stream did not block host until done; was already in an error state'; skipping this check. This only means that we won't check cudnn for out-of-bounds reads and writes. This message will only be printed once. 2023-10-05 16:02:22.318041: I tensorflow/stream_executor/cuda/cuda_driver.cc:739] failed to allocate 8B (8 bytes) from device: CUDA_ERROR_LAUNCH_TIMEOUT: the launch timed out and was terminated 2023-10-05 16:02:22.318214: E tensorflow/stream_executor/stream.cc:4476] INTERNAL: Failed to enqueue async memset operation: CUDA_ERROR_LAUNCH_TIMEOUT: the launch timed out and was terminated 2023-10-05 16:02:22.318464: E tensorflow/stream_executor/cuda/cuda_driver.cc:1163] failed to enqueue async memcpy from device to host: CUDA_ERROR_LAUNCH_TIMEOUT: the launch timed out and was terminated; host dst: 0x3d7d3db5c0; GPU src: (nil); size: 8=0x8 2023-10-05 16:02:22.319026: I tensorflow/stream_executor/stream.cc:4442] [stream=000002328328CF00,impl=000002329929D2E0] INTERNAL: stream did not block host until done; was already in an error state 2023-10-05 16:02:22.319321: E tensorflow/stream_executor/stream.cc:4476] INTERNAL: Failed to enqueue async memset operation: CUDA_ERROR_LAUNCH_TIMEOUT: the launch timed out and was terminated 2023-10-05 16:02:22.319515: E tensorflow/stream_executor/stream.cc:4476] INTERNAL: Failed to enqueue async memset operation: CUDA_ERROR_LAUNCH_TIMEOUT: the launch timed out and was terminated 2023-10-05 16:02:22.319700: F tensorflow/stream_executor/cuda/cuda_dnn.cc:215] Check failed: status == CUDNN_STATUS_SUCCESS (7 vs. 0)Failed to set cuDNN stream.

詳細

• 以下のコードでGPUの認識はされていることが確認できました。

python

1from tensorflow.python.client import device_lib 2import tensorflow as tf 3device_lib.list_local_devices()4# Created device /device:GPU:0 with 45894 MB memory: -> device: 0, name: NVIDIA RTX A6000, pci bus id: 0000:09:00.0, compute capability: 8.65 6with tf.device('/gpu:0'):7 a = tf.Variable([[1., 2.]])8 b = tf.constant([[3.], [4.]])9 print(tf.matmul(a, b))10# tf.Tensor([[11.]], shape=(1, 1), dtype=float32)

• nvidia-smiコマンドでGPUの動向を調査

・学習を走らせる前

・実行後

・一度この表示になると、再起動するまでGPUはOSにもTFにも認識されなくなります

軽量のCNNでMNISTの分類タスクを走らせると、バッチサイズ50くらいまでは動作します

試したこと

・GPUドライバーの更新・再インストール
・CUDA・cuDNNの再インストール
・システムファイルチェッカー
・システムファイルのクリーンアップ
・windowsメモリ診断
・ドライブのエラーチェック（chkdsk）

追加情報

Pytorchでも試してみましたが、こちらもエラーが発生しました。

RuntimeError: CUDA error: unknown error CUDA kernel errors might be asynchronously reported at some other API call, so the stack trace below might be incorrect. For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

お願い

原因がわからず長いこと悩まされています。
初心者なので些細なアドバイスでも嬉しいです。よろしくお願いします。