import tensorflow as tf

      if tf.DeviceSpec.from_string(device.name).device_type == "TPU":
        return (32,)
    return (16, 32)

  def _force_device_sync(self):
    """Shamelessly copied from `resnet50_test.py`."""
    tf.constant(1.).cpu()

  def _report(self, label, start, num_iters, device, batch_size, data_format):
    avg_time = (time.time() - start) / num_iters
    dev = tf.DeviceSpec.from_string(device).device_type.lower()
    name = "%s_%s_batch_%d_%s" % (label, dev, batch_size, data_format)
    extras = {"examples_per_sec": batch_size / avg_time}
    self.report_benchmark(
        iters=num_iters, wall_time=avg_time, name=name, extras=extras)

  def _benchmark_eager_apply(self,
                             label,
                             device_and_format,
                             defun=False,
