An example where prediction runs only on CPU and not on GPU

prediction_example.py

# Example from https://www.tensorflow.org/tutorials/keras/classification
import time
import tensorflow as tf
from tensorflow import keras
tf.debugging.set_log_device_placement(True)
import numpy as np

print(tf.__version__)

print("######################################### List physical GPUs")
print(tf.config.experimental.list_physical_devices('GPU'))


print("####################################### Traini model")
fashion_mnist = keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dense(10)
])


model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.fit(train_images, train_labels, epochs=10)

test_loss, test_acc = model.evaluate(test_images,  test_labels, verbose=2)

probability_model = tf.keras.Sequential([model,
                                         tf.keras.layers.Softmax()])

for i in range(2):
  print(f"################### [{i}] PREDICTION STARTED")
  time_start = time.perf_counter()
  probability_model.predict(test_images)
  print(f"################### [{i}] PREDICTION FINISHED in {(time.perf_counter() - time_start) * 1000} ms")