Начиная с версии 2 Tensorflow для расчетов по умолчанию использует видеокарту (если она настроена корректно).

Подобное бывает не всегда удобным, так как GPU раскрывает себя при расчетах больших пакетов на рекуррентных, сверточных или более сложных (например с использованием BERT, сверточных и полносвязных слоев) моделях. В тоже время при расчетах небольших моделей прямого распространения видеокарта оказывается медленнее CPU из-за задержек при копировании данных в и из видеопамяти.

Способы вручную указать устройство, которое надо использовать для расчетов.

Переменная окружения CUDA_VISIBLE_DEVICES

import os

...

# Использовать только процессор.
os.environ["CUDA_VISIBLE_DEVICES"]="-1"

Документация.

Назначение видимых устройств

import tensorflow as tf

...

physical_devices = tf.config.list_physical_devices()
# Использовать только процессор.
tf.config.set_visible_devices(physical_devices[2:])

Документация про:

• получение списка поддерживаемых устройств;
• установку видимых устройств;
• то, почему каждое устройство имеет клона с префиксом XLA_.

Контекстный менеджер устройства

import tensorflow as tf

...

with tf.device('/CPU:0'):
    # фиттинг модели