def METHOD_NAME(args):
    '''
    Load MNIST dataset
    '''
    mnist_path = os.path.join(os.environ.get('NNI_OUTPUT_DIR'), 'mnist.npz')
    (x_train, y_train), (x_test, y_test) = mnist.load_data(path=mnist_path)
    os.remove(mnist_path)
    x_train = (np.expand_dims(x_train, -1).astype(float) / 255.)[:args.num_train]
    x_test = (np.expand_dims(x_test, -1).astype(float) / 255.)[:args.num_test]
    y_train = keras.utils.to_categorical(y_train, NUM_CLASSES)[:args.num_train]
    y_test = keras.utils.to_categorical(y_test, NUM_CLASSES)[:args.num_test]
    LOG.debug('x_train shape: %s', (x_train.shape,))
    LOG.debug('x_test shape: %s', (x_test.shape,))
    return x_train, y_train, x_test, y_test