Introduction

Keras uses TensorFlow's device management, so choosing a specific GPU is really a TensorFlow configuration task. The safest approach is to limit visible devices before the model is created, which makes Keras behave as if only the selected GPU exists.

The Simplest Option: CUDA_VISIBLE_DEVICES

If you want the process to use only one GPU, set the environment variable before Python starts:

CUDA_VISIBLE_DEVICES=1 python train.py

Inside the process, TensorFlow now sees physical GPU 1 as logical GPU 0. This is often the easiest solution for scripts, notebooks launched from a shell, or job schedulers.

You can confirm what TensorFlow sees:

1import tensorflow as tf
2
3print(tf.config.list_physical_devices("GPU"))
4print(tf.config.list_logical_devices("GPU"))

Selecting the GPU in Code

If you prefer to choose the GPU programmatically, do it before creating tensors or models. Once TensorFlow initializes the runtime, device visibility changes may fail.

1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4if not gpus:
5    raise RuntimeError("No GPU found")
6
7target_gpu = gpus[1]
8tf.config.set_visible_devices(target_gpu, "GPU")
9tf.config.experimental.set_memory_growth(target_gpu, True)
10
11print(tf.config.list_logical_devices("GPU"))

After this, Keras uses only that selected device:

1model = tf.keras.Sequential([
2    tf.keras.layers.Input(shape=(10,)),
3    tf.keras.layers.Dense(32, activation="relu"),
4    tf.keras.layers.Dense(1)
5])
6
7model.compile(optimizer="adam", loss="mse")

A Complete Training Example

1import numpy as np
2import tensorflow as tf
3
4gpus = tf.config.list_physical_devices("GPU")
5tf.config.set_visible_devices(gpus[0], "GPU")
6tf.config.experimental.set_memory_growth(gpus[0], True)
7
8x = np.random.randn(128, 10).astype("float32")
9y = np.random.randn(128, 1).astype("float32")
10
11model = tf.keras.Sequential([
12    tf.keras.layers.Input(shape=(10,)),
13    tf.keras.layers.Dense(64, activation="relu"),
14    tf.keras.layers.Dense(1)
15])
16
17model.compile(optimizer="adam", loss="mse")
18model.fit(x, y, epochs=2, batch_size=16, verbose=0)

If only one GPU is visible to TensorFlow, Keras trains on that GPU automatically.

Verifying the Device Placement

If you want confirmation, enable device-placement logging:

tf.debugging.set_log_device_placement(True)

You can also inspect GPU utilization with external tools such as:

nvidia-smi

That is often the quickest real-world check.

When with tf.device(...) Helps

You may also see examples that wrap model creation or specific operations in a device context:

1with tf.device("/GPU:0"):
2    model = tf.keras.Sequential([
3        tf.keras.layers.Input(shape=(10,)),
4        tf.keras.layers.Dense(32, activation="relu"),
5        tf.keras.layers.Dense(1)
6    ])

This can be useful for placing a known block of work, but it is not usually the best primary method for selecting one GPU on a multi-GPU machine. In practice, limiting visible devices is clearer because it affects the whole process consistently, including dataset pipelines, layer creation, and training steps.

Memory Growth and Multi-GPU Machines

On multi-GPU systems, TensorFlow may try to reserve a large amount of GPU memory on visible devices. Enabling memory growth can reduce contention when several jobs share the same machine.

If you need to use more than one GPU intentionally, do not hide devices. Instead, look at distribution strategies such as tf.distribute.MirroredStrategy. But if the goal is one specific GPU, limiting visibility is simpler and usually more predictable.

Common Pitfalls

The most common mistake is calling set_visible_devices after TensorFlow has already initialized the GPU runtime by creating tensors, models, or datasets. Device configuration must happen first.

Another issue is forgetting that CUDA_VISIBLE_DEVICES=1 renumbers the visible device list inside the process. TensorFlow may report that single visible GPU as logical GPU 0, which is expected.

A third pitfall is assuming Keras has a separate GPU-selection API. It does not. Keras follows TensorFlow's device visibility and placement rules.

Finally, if TensorFlow cannot see any GPUs at all, the problem is not Keras configuration. Check the installed TensorFlow build, CUDA support, drivers, and container runtime first.

Summary

• Keras uses TensorFlow device configuration, not its own separate GPU selector.
• The simplest way to target one GPU is CUDA_VISIBLE_DEVICES.
• In-code selection works with tf.config.set_visible_devices if done before runtime initialization.
• Verify placement with TensorFlow logging or nvidia-smi.
• If no GPUs are visible, fix the TensorFlow and driver environment before debugging model code.