import tensorflow as tf

gpus = tf.config.list_physical_devices("GPU")
if not gpus:
    raise SystemExit("No TensorFlow GPU devices found")

for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

for gpu in gpus:
    enabled = tf.config.experimental.get_memory_growth(gpu)
    print(f"{gpu.name} memory_growth={enabled}")

logical_gpus = tf.config.list_logical_devices("GPU")
print(f"{len(gpus)} Physical GPUs, {len(logical_gpus)} Logical GPUs")

# Continue with model, dataset, and training code here.