def test_center_batch_shape(self):
    """Test that the shape of the target for a batch is correct."""
    def graph_fn():
      box_batch = [
          tf.constant([self._box_center, self._box_lower_left]),
          tf.constant([self._box_center]),
          tf.constant([self._box_center_small]),
      ]
      classes = [
          tf.one_hot([0, 1], depth=4),
          tf.one_hot([2], depth=4),
          tf.one_hot([3], depth=4),
      ]
      assigner = targetassigner.CenterNetCenterHeatmapTargetAssigner(4)
      targets = assigner.assign_center_targets_from_boxes(80, 80, box_batch,
                                                          classes)
      return targets
    targets = self.execute(graph_fn, [])
    self.assertEqual((3, 20, 20, 4), targets.shape)
Test that the shape of the target for a batch is correct.
inference_api/src/main/object_detection/core/target_assigner_test.py