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

[2son1a]

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