def unflatten_images(input_batch, depth, height, width):
    """
    Take a batch of images and unflatten into a DxHxW grid.
    Nearly an inverse of `flatten_images`. (`flatten_images` assumes a list of tensors, not a tensor.)
    
    Args:
    * input_batch: a tensor of dtype=float and shape (bsz, d*h*w).
    * depth: int
    * height: int
    * width: int
    """
    return input_batch.view(input_batch.shape[0], depth, height, width)