import tensorflow as tf

    input_partition_dims = None
    num_cores_per_replica = None

    if params.use_tpu:
      tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
          params.platform.tpu,
          zone=params.platform.tpu_zone,
          project=params.platform.gcp_project)
      tpu_grpc_url = tpu_cluster_resolver.get_master()
      tf.Session.reset(tpu_grpc_url)

      # If the input image is transposed (from NHWC to HWCN), the partition
      # dimensions also need to be transposed the same way.
      def _maybe_transpose(input_partition_dims):
        if input_partition_dims and params.train.transpose_input:
          return [input_partition_dims[i] for i in [1, 2, 3, 0]]
        else:
          return input_partition_dims
