如何在Tensorflow 2.0中使用Google Colab TPU将检查点保存到Google Cloud Storage？

Question

1. Tensorflow 1.X中的脚本，但在 Tensorflow 2.0中仍然无法使用

2。我该如何将其迁移到 Tensorflow 2.0

import datetime
import json
import os
import pprint
import random
import string
import sys
import tensorflow as tf
import tensorflow_io as tfio

assert 'COLAB_TPU_ADDR' in os.environ, 'ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!'
TPU_ADDRESS = 'grpc://' + os.environ['COLAB_TPU_ADDR']
print('TPU address is', TPU_ADDRESS)

from google.colab import auth
auth.authenticate_user()

with tf.compat.v1.Session(TPU_ADDRESS) as session:
  print('TPU devices:')
  pprint.pprint(session.list_devices())

  # Upload credentials to TPU.
  with open('/content/adc.json', 'r') as f:
    auth_info = json.load(f)
  tf.contrib.cloud.configure_gcs(session, credentials=auth_info)
#   Now credentials are set for all future sessions on this TPU.