code: def get_generator_loss(
      self,
      batch: tensor_dict_type,
      forward_results: tensor_dict_type,
      *,
      step: Optional[int],
      last_layer: nn.Parameter,
      cond: Optional[Tensor] = None,
      nll_weights: Optional[float] = None,
  ) -> CustomTrainStepLoss:
      inputs = batch[INPUT_KEY].contiguous()
      reconstructions = forward_results[PREDICTIONS_KEY].contiguous()
      loss_items = {}
      # recon & nll & kl loss
      recon_loss = torch.abs(inputs - reconstructions)
      with torch.no_grad():
          loss_items["l1"] = recon_loss.mean().item()
      if self.perceptual_weight > 0:
          p_loss = self.perceptual_loss(inputs, reconstructions)
          p_loss = self.perceptual_weight * p_loss
          recon_loss = recon_loss + p_loss
          with torch.no_grad():
              loss_items["perceptual"] = p_loss.mean().item()
      with torch.no_grad():
          loss_items["recon"] = recon_loss.mean().item()
      nll_loss = recon_loss / torch.exp(self.log_var) + self.log_var
      weighted_nll_loss = nll_loss
      if nll_weights is not None:
          weighted_nll_loss = nll_weights * nll_loss
      weighted_nll_loss = (
          torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
      )
      nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
      kl_loss = forward_results[GaussianDistribution.key].kl()
      kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
      loss = weighted_nll_loss + self.kl_weight * kl_loss
      loss_items["kl"] = self.kl_weight * kl_loss.item()
      # check discriminator start
      if step is not None and step < self.d_loss_start_step:
          return CustomTrainStepLoss(loss, loss_items)
      g_loss = self.g_loss(
          nll_loss, last_layer, loss_items, reconstructions, cond
      )
      loss = loss + g_loss
      return CustomTrainStepLoss(loss, loss_items)
  
imports: import torch
  import numpy
  import torch.nn
  from torch import Tensor
  from typing import Tuple
  from typing import Optional
  from cftool.types import tensor_dict_type
  from .common import IAutoEncoder
  from .common import AutoEncoderInit
  from .common import AutoEncoderLPIPSWithDiscriminator
  from ...schemas import CustomTrainStepLoss
  from ....schema import IDLModel
  from ....constants import INPUT_KEY
  from ....constants import PREDICTIONS_KEY
label_desc: Use image node_id 2 for calling the AutoEncoderKLLoss obj's underlying member method code with example usage: obj.get_generator_loss(batch, forward_results) and returns: CustomTrainStepLoss, CustomTrainStepLoss