Prueba tensorflow GPU

import tensorflow as tf

# Verifica si TensorFlow fue construido con soporte para CUDA (GPU)
print("¿TensorFlow construido con CUDA?:", tf.test.is_built_with_cuda())

# Verifica si TensorFlow puede acceder a una GPU
print("¿TensorFlow puede acceder a una GPU?:", tf.test.is_gpu_available())

# Lista las GPUs disponibles y sus capacidades
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    print("GPUs disponibles:")
    for gpu in gpus:
        print(gpu)
else:
    print("No se encontraron GPUs disponibles.")

# Verifica la versión de CUDA y cuDNN si están disponibles
if tf.test.is_built_with_cuda():
    print("Versión de CUDA:", tf.sysconfig.get_build_info()['cuda_version'])
    print("Versión de cuDNN:", tf.sysconfig.get_build_info()['cudnn_version'])

# Crea un modelo simple con una capa convolucional
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# Compila el modelo
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# Crea datos de entrada y salida falsos
import numpy as np
x = np.random.random((100, 28, 28, 1))
y = np.random.randint(0, 10, size=(100,))

# Entrena el modelo
model.fit(x, y, epochs=1)