def process_batch(images):
    assert images.shape == (32, 28, 28, 1), "Wrong image shape!"
    return images * 2.0

        __
        __