Keras, AIX360(LIME) - ValueError: the input array must be have a shape == (.., ..,[ ..,] 3)), got (28, 28, 1)

I am trying to make some program using a model from Keras and then explain it with Lime explainer from AIX360 (which is just wrapper for the original LIME). All the data are MNIST grayscale digits. But in my case, I am not able to explain the instance because I can't figure out what to feed to the explainers.

My code:

!pip install aix360
!pip install tensorflow==2.2.0

from __future__ import print_function
import warnings
# Supress jupyter warnings if required for cleaner output
warnings.simplefilter('ignore')

import numpy as np
import pandas as pd

import keras
import keras.layers

from keras.utils.np_utils import to_categorical # convert to one-hot-encoding
from keras.models import Sequential # Sequeantial layer addition

from aix360.algorithms.lime import LimeImageExplainer

print('Using keras:', keras.__version__)

# Load dataset
from keras.datasets import mnist
# Tuple of Numpy arrays: (x_train, y_train), (x_test, y_test).
(train, train_labels), (test, test_labels) = mnist.load_data()

# save input image dimensions
img_rows = train.shape[1]
img_cols = train.shape[2]

# Get classes and number of values
value_counts = pd.value_counts(train_labels).sort_index()
num_classes = value_counts.count()

train = train/255
test = test/255

train = train.reshape(train.shape[0], img_rows, img_cols, 1)
test = test.reshape(test.shape[0], img_rows, img_cols, 1)

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(img_rows, img_cols,1)),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dense(num_classes)
])

model.compile(loss='sparse_categorical_crossentropy',
      optimizer='adam',
      metrics=['accuracy'])

batch_size = 128
epochs = 1

model.fit(train, train_labels,
          batch_size=batch_size,
          epochs=epochs,
          verbose=1,
          validation_data=(test, test_labels))

score = model.evaluate(test, test_labels, verbose=0)

print('Test loss:', score[0])
print('Test accuracy:', score[1])

limeExplainer = LimeImageExplainer()

limeExplainer.explain_instance(test[0], model.predict_proba)

The last line is where the error is. Don't pay attention to how the model is trained, that is not a concern.

EDIT: edited code so it is hopefully runable in codelab (added second line)

EDIT2: To be complete: tensorflow 2.2.0 keras 2.4.3 aix360 0.2.0