OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed trying to use numpy()

I have some tfrecord files which I am loading into a TFRecordDataset and then trying to pass to dataset.map. There is a lot of data in this dataset, and I am just trying to extract an image and 2D bounding boxes from each record. I'm using the Waymo dataset and its api for extracting the data.

In order to extract this information using the Waymo API I have to convert the Tensor to a numpy array. I call my map function like this

dataset.map(read_tfrecord, num_parallel_calls=4)

the mapping function is here

@tf.function
def read_tfrecord(example):
  image, true_boxes = tf.py_function(read_tfrecord_np, [example], tf.float32)
  return image, true_boxes

My read_tfrecord_np function looks like

def read_tfrecord_np(example):
    frame = open_dataset.Frame()
    frame.ParseFromString(bytearray(example.numpy()))
    ...
    return front_image, true_boxes

Here is the stacktrace

tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: in converted code: train.py:66 read_tfrecord *
true_boxes, image = tf.py_function(read_tfrecord_np, [example], tf.float32)
/home/vagrant/anaconda3/envs/projf/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:539 iter
self._disallow_iteration() /home/vagrant/anaconda3/envs/projf/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:532 _disallow_iteration
self._disallow_when_autograph_enabled("iterating over tf.Tensor")
/home/vagrant/anaconda3/envs/projf/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:510 _disallow_when_autograph_enabled
" decorating it directly with @tf.function.".format(task)) OperatorNotAllowedInGraphError: iterating over tf.Tensor is not allowed: AutoGraph did not convert this function. Try decorating it directly with @tf.function.

It says "try decorating it directly with @tf.function" but I am doing that. What am I doing wrong here? I'm using tensorflow 2.1.0