def reconstruction_loss(reconstruction, orig):
    """
    Computes the reconstruction loss
        :param reconstruction: batch x im_size
        :param orig: batch x im_ht x im_wd
        :return rloss: The reconstruction loss
    """
    rloss = ((reconstruction - orig.view(orig.size(0),
                                         orig.size(1), -1)) ** 2).mean(-1)
    rloss = rloss.mean()
    return rloss