def METHOD_NAME(num_class, batch_size):
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train = np.expand_dims(x_train, axis=-1).astype("float32")
    x_test = np.expand_dims(x_test, axis=-1).astype("float32")
    x_train /= 255  # normalize to 0~1
    x_test /= 255
    y_train = np_utils.to_categorical(y_train, num_class)
    y_test = np_utils.to_categorical(y_test, num_class)
    train_ds = (
        tf.data.Dataset.from_tensor_slices((x_train, y_train))
        .repeat()
        .batch(batch_size, drop_remainder=True)
    )
    eval_ds = (
        tf.data.Dataset.from_tensor_slices((x_test, y_test))
        .repeat()
        .batch(batch_size, drop_remainder=True)
    )
    return train_ds, eval_ds