Tensorflow Lite Model Maker error: object has no attribute 'index_to_label'

Question

I'm usign Tensorflow Lite Model Maker to adapt a model to citrus_leaves dataset, but following the documentation I get this error:

AttributeError: 'PrefetchDataset' object has no attribute 'index_to_label'

I'm making this code on Google colab:

!pip install -q tflite-model-maker
!pip install tensorflow-datasets
!pip install tfds-nightly

import os
import numpy as np
import tensorflow as tf
assert tf.__version__.startswith('2') 
from tflite_model_maker import model_spec
from tflite_model_maker import image_classifier
from tflite_model_maker.config import ExportFormat
from tflite_model_maker.config import QuantizationConfig
from tflite_model_maker.image_classifier import DataLoader
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt

(test_data, train_data), info = tfds.load('citrus_leaves', split=["train[:20%]", "train[20%:]"], shuffle_files=True, with_info=True, as_supervised=True)

model = image_classifier.create(train_data) # Here the error happens

loss, accuracy = model.evaluate(test_data)

Here is the Link to dataset and here is the link to the docs that I was following.

I have no idea why this error is happening, I would be grateful if someone could help me. Thanks!