tfds errors with tfds.testing.mock_data

What I need help with / What I was wondering
How to use tfds.testing.mock_data correctly

What I've tried so far
This code on its own works fine:

import tensorflow_datasets as tfds
ds_builder = tfds.builder('mnist')
ds_builder.download_and_prepare()
train_ds = tfds.as_numpy(ds_builder.as_dataset(split='train', batch_size=-1))

When decorating that snippet with tfds.testing.mock_data I get errors:

with tfds.testing.mock_data(num_examples=8, data_dir='/workdisk/.tfds2/metadata'):
  import tensorflow_datasets as tfds
  ds_builder = tfds.builder('mnist')
  ds_builder.download_and_prepare()
  train_ds = tfds.as_numpy(ds_builder.as_dataset(split='train', batch_size=-1))

Error trace:

I1220 13:56:57.120099 140470378489664 mocking.py:162] Metadata NOT found for mnist at /workdisk/jax/FlaxExamples/flax/.tfds/metadata/mnist/3.0.1. Will use `MockPolicy.USE_CODE.`
I1220 13:56:57.121221 140470378489664 logging_logger.py:35] Constructing tf.data.Dataset mnist for split train, from /workdisk/jax/FlaxExamples/flax/.tfds/metadata/mnist/3.0.1
WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/tensorflow_datasets/core/dataset_builder.py:622: get_single_element (from tensorflow.python.data.experimental.ops.get_single_element) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Dataset.get_single_element()`.
W1220 13:56:57.204224 140470378489664 deprecation.py:341] From /usr/local/lib/python3.8/dist-packages/tensorflow_datasets/core/dataset_builder.py:622: get_single_element (from tensorflow.python.data.experimental.ops.get_single_element) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.data.Dataset.get_single_element()`.
terminate called after throwing an instance of 'std::length_error'
  what():  vector::reserve
Fatal Python error: Aborted

Environment information
Latest Nvidia dev docker image: nvidia/cuda:11.4.2-cudnn8-devel-ubuntu20.04

• Operating System: ubuntu20.04
• Python version: 3.8
• tensorflow 2.7.0
• tensorflow-datasets 4.4.0
• tensorflow-estimator 2.7.0
• tensorflow-io-gcs-filesystem 0.23.1
• tensorflow-metadata 1.5.0