.predict() runs only on CPU even though GPU is available

Since you already have a GPU, I assume that tf.test.is_gpu_available() returns True. You can use this piece of code to force TensorFlow to use a specific device-

with tf.device('/gpu:0'):
    // GPU stuff

This also works if you want to force it to use a CPU instead for some part of the code-

with tf.device('/cpu:0'):
    // CPU stuff

An addon which might be helpful while using tf.device(), you can use this function to list all the devices you have-

def get_available_devices():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos]

get_available_devices()

Though for the use-case you mentioned, I do not guarantee faster inferences with a GPU.

Sounds like you need to use a Distributed Strategy per the docs. Your code would then become something like the following:

tf.debugging.set_log_device_placement(True)
strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    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()]
    )
    probability_model.predict(test_images)

Per the documentation, The best practice for using multiple GPUs is to use tf.distribute.Strategy.

Your predict function is using GPU. And, I have recalculated timing on NVIDIA 1080 GTX with your code & it is taking 100 ms for inference.

Either reboot the system or check if GPU is getting utilised or not.

Here is the line of your code stating inference is run on GPU:

2020-09-01 06:19:15.885778: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op __inference_distributed_function_58022 in device /job:localhost/replica:0/task:0/device:GPU:0

Are you using the correct tensorflow package? It could help to uninstall tensorflow and install tensorflow-gpu instead.

For documentation see: https://www.tensorflow.org/install/gpu

Can you please call keras from tensorflow.python?

Like;

from tensorflow.python.keras.models import Sequential

Also, please check CUDA and CuDNN version. Cuda and CuDNN versions have to be compatible with tensorflow version. You can check from here. As your tensorflow version is 2.1, CUDA and CuDNN versions are have to be 10.1 & 7.6 respectively.