The code works when omniglot_split="train", but crashes when omniglot_split="test"

Hi!

I am trying to run the code below. It works when omniglot_split="train", but crashes when omniglot_split="test". Could you tell me what the problem might be?

The code:

```
from datasets import data_generators
import tensorflow as tf
from tensorflow_datasets.core.utils import gcs_utils
from torch.utils.data import Dataset
import torch

dataset_for_sampling = data_generators.OmniglotDatasetForSampling(
omniglot_split="test",
exemplars="single",
augment_images=False,
)

generator = data_generators.SeqGenerator(
dataset_for_sampling,
n_rare_classes=10,
n_common_classes=10,
n_holdout_classes=10,
zipf_exponent=1,
use_zipf_for_common_rare=False,
noise_scale=0.1,
preserve_ordering_every_n=None,
random_seed=13,
)

generator_factory = generator.get_fewshot_seq
output_signature = {
'example':
tf.TensorSpec(
shape=(9, 105, 105, 1), dtype=tf.dtypes.float32),
'label':
tf.TensorSpec(shape=(9,), dtype=tf.dtypes.int32),
'is_rare':
tf.TensorSpec(shape=(9,), dtype=tf.dtypes.int32)
}

dataset = tf.data.Dataset.from_generator(
generator_factory,
args=('common', 4, 2, 'unfixed', False, False),
output_signature=output_signature
)

for example in dataset:
print("example")
```

The error:

```
Traceback (most recent call last):

File "/opt/conda/lib/python3.10/site-packages/tensorflow/python/ops/script_ops.py", line 270, in __call__
ret = func(*args)

File "/opt/conda/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
return func(*args, **kwargs)

File "/opt/conda/lib/python3.10/site-packages/tensorflow/python/data/ops/from_generator_op.py", line 198, in generator_py_func
values = next(generator_state.get_iterator(iterator_id))

File "/app/icl/datasets/data_generators.py", line 787, in get_fewshot_seq
seq_examples = self._create_noisy_image_seq(

File "/app/icl/datasets/data_generators.py", line 276, in _create_noisy_image_seq
if isinstance(self.data[classes[0]], list):

KeyError: 136
```