title: Why does Google Colab give a "unknown device" error when I try to convert Keras MLP?
body: <p>I'm trying to train a simple MLP model on google colab using a TPU.
  However, when I try to convert the model with</p>
  
  <pre class="lang-py prettyprint-override"><code>from tensorflow.keras.models import Sequential
  from tensorflow.keras.layers import Dense
  from keras.constraints import NonNeg
  model = Sequential()
  model.add(Dense(57,input_shape=(57,)))
  model.add(Dense(60,kernel_constraint=NonNeg(),activation="relu"))
  model.add(Dense(100,kernel_constraint=NonNeg(),activation="relu"))
  model.add(Dense(50,kernel_constraint=NonNeg(),activation="relu"))
  model.add(Dense(3,activation="linear"))
  model.compile(optimizer='adam',loss='mse')
  
  TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
  tpu_model = tf.contrib.tpu.keras_to_tpu_model(
      model,
      strategy=tf.contrib.tpu.TPUDistributionStrategy(
          tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))
  </code></pre>
  
  <p>gives the error <code>InvalidArgumentError: /job:localhost/replica:0/task:0/device:TPU_SYSTEM:0 unknown device.</code> <a href="https://pastebin.com/YVUeWTrU" rel="nofollow noreferrer">full error here</a>.</p>
  
  <p>This doesn't make sense to me, because I copied the model conversion code straight from the google tutorial. When I run the test code:</p>
  
  <pre class="lang-py prettyprint-override"><code>import os
  import pprint
  import tensorflow as tf
  
  if 'COLAB_TPU_ADDR' not in os.environ:
    print('ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!')
  else:
    tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
    print ('TPU address is', tpu_address)
  
    with tf.Session(tpu_address) as session:
      devices = session.list_devices()
  
    print('TPU devices:')
    pprint.pprint(devices)
  </code></pre>
  
  <p>I get the expected result:</p>
  
  <pre><code>TPU address is grpc://10.0.203.10:8470
  TPU devices:
  [_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:CPU:0, CPU, -1, 1941375595625340814),
   _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 9079881000847066378),
   _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 6922694346479333534),
   _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 14324637633413341896),
   _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 3528106575831937158),
   _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 13852141601322651612),
   _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 10344791506504172772),
   _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 16666353711371098164),
   _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 3428083526573573796),
   _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 8632908473312514763),
   _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 17179869184, 9715206562754100387)]
  </code></pre>
  
  <p>Since that works, I don't understand why the model conversion does not.</p>