Is it possible to only load part of a TensorFlow dataset?

Question

I have a notebook in Google Colab with the following code:

batch_size = 64
dataset_name = \'coco/2017_panoptic\'

tfds_dataset, tfds_info = tfds.load(
    dat