NVIDIA cuDNN Runtime Libraries for CUDA 12

import tensorflow as tf
import os

# Ensure TensorFlow doesn't pre-allocate all GPU memory
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# Check if TensorFlow can detect and use GPUs
gpus = tf.config.list_physical_devices('GPU')

if gpus:
    print(f"TensorFlow detected the following GPUs: {gpus}")
    try:
        # Limit GPU memory growth to avoid allocating all memory at once (alternative to env var)
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        print("GPU memory growth set to True.")
    except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print(f"Error setting memory growth: {e}")
    print(f"TensorFlow is built with CUDA: {tf.test.is_built_with_cuda()}")
    # TensorFlow's built-in cuDNN version (indicates what TF was compiled with)
    print(f"TensorFlow's built-in cuDNN version: {tf.sysconfig.get_build_info().get('CUDNN_VERSION', 'N/A')}")

    # A small operation to trigger GPU usage if available
    try:
        with tf.device('/GPU:0'):
            a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
            b = tf.constant([[1.0, 1.0], [1.0, 1.0]])
            c = tf.matmul(a, b)
            print(f"Simple matrix multiplication on GPU: {c.numpy()}")
    except RuntimeError as e:
        print(f"Could not run on GPU: {e}. Running on CPU instead.")
        a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
        b = tf.constant([[1.0, 1.0], [1.0, 1.0]])
        c = tf.matmul(a, b)
        print(f"Simple matrix multiplication on CPU: {c.numpy()}")
else:
    print("TensorFlow did not detect any GPUs. Please ensure CUDA and cuDNN are correctly installed and configured.")