class_docstr: ## Reconstruction loss
  
  $$L_{Rec} = \sum_{n=1}^N p_n \mathcal{L}(y, \hat{y}_n)$$
  
  $\mathcal{L}$ is the normal loss function between target $y$ and prediction $\hat{y}_n$.
class_docstr_tok: ["#","#","Reconstruction","loss","$","$","L_","{","Rec","}","=","\\sum_","{","n=1","}","^N","p_n","\\mathcal","{","L","}","(","y",",","\\hat","{","y","}","_n",")","$","$","$","\\mathcal","{","L","}","$","is","the","normal","loss","function","between","target","$","y","$","and","prediction","$","\\hat","{","y","}","_n","$","."]
docstr: * `p` is $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
  * `y_hat` is $\hat{y}_1 \dots \hat{y}_N$ in a tensor of shape `[N, batch_size, ...]`
  * `y` is the target of shape `[batch_size, ...]`
docstr_tok: ["*","`","p","`","is","$","p_1","\\dots","p_N","$","in","a","tensor","of","shape","`","[","N",",","batch_size","]","`","*","`","y_hat","`","is","$","\\hat","{","y","}","_1","\\dots","\\hat","{","y","}","_N","$","in","a","tensor","of","shape","`","[","N",",","batch_size",",","...","]","`","*","`","y","`","is","the","target","of","shape","`","[","batch_size",",","...","]","`"]
code: def forward(
      self, p: torch.Tensor, y_hat: torch.Tensor, y: torch.Tensor
  ):
      """
      * `p` is $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
      * `y_hat` is $\hat{y}_1 \dots \hat{y}_N$ in a tensor of shape `[N, batch_size, ...]`
      * `y` is the target of shape `[batch_size, ...]`
      """
  
      # The total $\sum_{n=1}^N p_n \mathcal{L}(y, \hat{y}_n)$
      total_loss = p.new_tensor(0.0)
      # Iterate upto $N$
      for n in range(p.shape[0]):
          # $p_n \mathcal{L}(y, \hat{y}_n)$ for each sample and the mean of them
          loss = (p[n] * self.loss_func(y_hat[n], y)).mean()
          # Add to total loss
          total_loss = total_loss + loss
  
      #
      return total_loss
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/adaptive_computation/ponder_net/__init__.py
imports: from typing import Tuple
  import torch
  from torch import nn
  from labml_helpers.module import Module
label_desc: Use image node_id 2 for calling the ReconstructionLoss obj's underlying member method code with example usage: obj.forward(p, y_hat, y) and returns: total_loss