def test_tpu_only_test(self):
    """Set up tests only to run with TPU."""
    if self.has_tpu():
      def graph_fn(tensora, tensorb):
        return tensora + tensorb, tensora - tensorb
      tensora_np = np.ones(20)
      tensorb_np = tensora_np * 2
      output = self.execute_tpu(graph_fn, [tensora_np, tensorb_np])
      self.assertAllClose(output[0], tensora_np + tensorb_np)
      self.assertAllClose(output[1], tensora_np - tensorb_np)
Set up tests only to run with TPU.
inference_api/src/main/object_detection/utils/test_case_test.py