import tensorflow as tf

        # Parse 'None' into None.
        input_partition_dims = [
            None if x == 'None' else _maybe_transpose(x)
            for x in input_partition_dims
        ]
    else:
      tpu_cluster_resolver = None

    # Sets up config for TPUEstimator.
    tpu_config = tf.contrib.tpu.TPUConfig(
        params.train.iterations_per_loop,
        num_cores_per_replica=num_cores_per_replica,
        input_partition_dims=input_partition_dims,
        per_host_input_for_training=tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2  # pylint: disable=line-too-long
    )

    run_config = tf.contrib.tpu.RunConfig(
        cluster=tpu_cluster_resolver,
