code: def forward(
      self,
      forward_results: tensor_dict_type,
      batch: tensor_dict_type,
      *,
      reduction: str = "mean",
      gather: bool = True,
  ) -> losses_type:
      # reconstruction loss
      original = batch[INPUT_KEY]
      reconstruction = forward_results[PREDICTIONS_KEY]
      if self.loss_type == "l2":
          recon = F.mse_loss(
              reconstruction, original, reduction=reduction
          )
      elif self.loss_type == "l1":
          recon = F.l1_loss(
              reconstruction, original, reduction=reduction
          )
      else:
          raise ValueError(
              f"unrecognized loss_type '{self.loss_type}' occurred"
          )
      # vq & commit loss
      z_e = forward_results["z_e"]
      z_q_g = forward_results["z_q_g"]
      vq_loss = F.mse_loss(z_q_g, z_e.detach(), reduction=reduction)
      commit_loss = F.mse_loss(z_e, z_q_g.detach(), reduction=reduction)
      codebook_loss = (
          self.lb_vq * vq_loss + self.lb_commit * commit_loss
      )
      losses = {
          self.loss_type: recon,
          "commit": commit_loss,
          "codebook": codebook_loss,
      }
      if not gather:
          return losses
      # gather
      loss = self.lb_recon * recon + codebook_loss
      losses[LOSS_KEY] = loss
      return losses
  
file_path: carefree-learn/cflearn/losses/vae.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from typing import Optional
  from cftool.types import tensor_dict_type
  from ..types import losses_type
  from ..schema import ILoss
  from ..schema import TrainerState
  from ..constants import LOSS_KEY
  from ..constants import INPUT_KEY
  from ..constants import PREDICTIONS_KEY
label_desc: Use image node_id 2 for calling the VQVAELoss obj's underlying member method code with example usage: obj.forward(forward_results, batch) and returns: losses, losses