title: .predict() runs only on CPU even though GPU is available
body: <p>I used <a href="https://www.tensorflow.org/tutorials/keras/classification" rel="noreferrer">this script to train a model &amp; predict</a> on a machine with GPU installed and enabled and it seems that it's using only the CPU in the prediction stage.</p>
  <p>The device placement log I'm seeing during the <code>.predict()</code> part is the following:</p>
  <pre><code>2020-09-01 06:08:19.085400: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op RangeDataset in device /job:localhost/replica:0/task:0/device:CPU:0
  2020-09-01 06:08:19.085617: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op RepeatDataset in device /job:localhost/replica:0/task:0/device:CPU:0
  2020-09-01 06:08:19.089558: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op MapDataset in device /job:localhost/replica:0/task:0/device:CPU:0
  2020-09-01 06:08:19.090003: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op PrefetchDataset in device /job:localhost/replica:0/task:0/device:CPU:0
  2020-09-01 06:08:19.097064: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op FlatMapDataset in device /job:localhost/replica:0/task:0/device:CPU:0
  2020-09-01 06:08:19.097647: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op TensorDataset in device /job:localhost/replica:0/task:0/device:CPU:0
  2020-09-01 06:08:19.097802: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op RepeatDataset in device /job:localhost/replica:0/task:0/device:CPU:0
  2020-09-01 06:08:19.097957: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op ZipDataset in device /job:localhost/replica:0/task:0/device:CPU:0
  2020-09-01 06:08:19.101284: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op ParallelMapDataset in device /job:localhost/replica:0/task:0/device:CPU:0
  2020-09-01 06:08:19.101865: I tensorflow/core/common_runtime/eager/execute.cc:573] Executing op ModelDataset in device /job:localhost/replica:0/task:0/device:CPU:0
  </code></pre>
  <p>even though that when I run:</p>
  <pre><code>print(tf.config.experimental.list_physical_devices('GPU'))
  </code></pre>
  <p>I receive:</p>
  <pre><code>[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:2', device_type='GPU')]
  </code></pre>
  <p>The code I used <a href="https://gist.github.com/georgehdd/4bfd9d6110eefd26821bdcd5bce8c08c" rel="noreferrer">can be found here</a>. The <a href="https://gist.github.com/georgehdd/44cc8900024985a7f3ad07de5b969378" rel="noreferrer">full output logs can be seen here</a>.</p>
  <p>More context:<br />
  Python: <code>3.7.7</code><br />
  Tensorflow: <code>2.1.0</code><br />
  GPU: <code>Nvidia Tesla V100-PCIE-16GB</code><br />
  CPU: <code>Intel Xeon Gold 5218 CPU @ 2.30GHz</code><br />
  RAM: <code>394851272 KB</code><br />
  OS: <code>Linux</code></p>