scope.py

with strategy.scope():
  """
  This essentailly takes our model and makes it 
  compatible to train on a TPU.
  """
  model = tf.keras.models.Sequential()
  model.add(tf.keras.layers.Dense(512, input_size=[784,], activation='relu'))
  model.add(tf.keras.layers.Dense(256, activation='relu'))
  model.add(tf.keras.layers.Dense(128, activation='relu'))
  model.add(tf.keras.layers.Dense(64, activation='relu'))
  model.add(tf.keras.layers.Dense(10, activation='softmax'))

  """
  Compiling the model using the RMSProp optimizer 
  and Sparse Categorical Crossentropy loss
  """
  model.compile(
    optimizer=tf.train.RMSPropOptimizer(learning_rate=1e-2),
    loss=tf.keras.losses.sparse_categorical_crossentropy,
    metrics=['sparse_categorical_accuracy']
  )

  model.summary()