message: small fix for the autoclustering_xla.ipynb
diff: " model.set_weights(initial_weights)\n",
  "\n",
  "warmup(model, x_train, y_train, x_test, y_test)\n",
  - "%time train_model(model, x_train, y_train, x_test, y_test)\n",
  + "train_model(model, x_train, y_train, x_test, y_test)\n",
  "\n",
  "scores = model.evaluate(x_test, y_test, verbose=1)\n",
  "print('Test loss:', scores[0])\n",
  