def load_data(label=2):
    'Loads the MNIST dataset and preprocesses it: scales to [0.0, 1.0]\n    range, resizes the images from (28, 28) to (32, 32) and filters the\n    dataset to keep images of just one class.\n\n    Parameters\n    ----------\n    label : {2, 0, 1, 3, 4, 5, 6, 7, 8, 9}, default 2\n        The class of images to train and test on.\n\n    Returns\n    -------\n    (x_train, x_test) : tuple of ndarrays\n        NumPy arrays of training and testing images.\n    '
    ((x_train, y_train), (x_test, y_test)) = tf.keras.datasets.mnist.load_data()
    ids_train = np.where((y_train == label))
    ids_test = np.where((y_test == label))
    x_train = x_train[ids_train]
    x_test = x_test[ids_test]
    x_train = (x_train.astype(np.float32) / 255.0)
    x_test = (x_test.astype(np.float32) / 255.0)
    x_train = x_train[(..., None)]
    x_test = x_test[(..., None)]
    x_train = tf.image.resize(x_train, (32, 32))
    x_test = tf.image.resize(x_test, (32, 32))
    return (x_train, x_test)
Loads the MNIST dataset and preprocesses it: scales to [0.0, 1.0]
range, resizes the images from (28, 28) to (32, 32) and filters the
dataset to keep images of just one class.

Parameters
----------
label : {2, 0, 1, 3, 4, 5, 6, 7, 8, 9}, default 2
    The class of images to train and test on.

Returns
-------
(x_train, x_test) : tuple of ndarrays
    NumPy arrays of training and testing images.
def load_data(label=2):
    'Loads the MNIST dataset and preprocesses it: scales to [0.0, 1.0]\n    range, resizes the images from (28, 28) to (32, 32) and filters the\n    dataset to keep images of just one class.\n\n    Parameters\n    ----------\n    label : {2, 0, 1, 3, 4, 5, 6, 7, 8, 9}, default 2\n        The class of images to train and test on.\n\n    Returns\n    -------\n    (x_train, x_test) : tuple of ndarrays\n        NumPy arrays of training and testing images.\n    '
    ((x_train, y_train), (x_test, y_test)) = tf.keras.datasets.mnist.load_data()
    ids_train = np.where((y_train == label))
    ids_test = np.where((y_test == label))
    x_train = x_train[ids_train]
    x_test = x_test[ids_test]
    x_train = (x_train.astype(np.float32) / 255.0)
    x_test = (x_test.astype(np.float32) / 255.0)
    x_train = x_train[(..., None)]
    x_test = x_test[(..., None)]
    x_train = tf.image.resize(x_train, (32, 32))
    x_test = tf.image.resize(x_test, (32, 32))
    return (x_train, x_test)