code: def compute_loss(self, batch, return_output=False):
      # first compute our regular forced decoding loss
      token_loss, model_output = super().compute_loss(
          batch, return_output=True
      )
      notnull = batch.label_vec.ne(self.NULL_IDX)
      num_tokens = notnull.long().sum().item()
  
      encoder_states = model_output[2]
      ctx_know_attn = encoder_states[2]
  
      if self.knowledge_alpha == 0.0:
          loss = token_loss
      else:
          _, know_pred = ctx_know_attn.max(1)
          know_acc = (know_pred == batch.cs_ids).float().sum().item()
          know_chance = (
              batch.ck_mask.sum(1).float().reciprocal().sum().item()
          )
          self.metrics["know_chance"] += know_chance
          self.metrics["bsz"] += batch.text_vec.size(0)
          self.metrics["know_acc"] += know_acc
          know_loss = th.nn.functional.cross_entropy(
              ctx_know_attn, batch.cs_ids, reduction="mean"
          )
          self.metrics[
              "know_loss"
          ] += know_loss.item() * batch.text_vec.size(0)
          # in the original paper the loss was scaled by num_tokens for both
          # know_loss and token_loss
          know_loss /= num_tokens
          loss = (
              1 - self.knowledge_alpha
          ) * token_loss + self.knowledge_alpha * know_loss
      if return_output:
          return (loss, model_output)
      else:
          return loss
  
imports: from typing import Optional
  from parlai.core.params import ParlaiParser
  from parlai.core.opt import Opt
  from itertools import chain
  from functools import lru_cache
  import torch
  import numpy
  from parlai.utils.torch import padded_tensor
  from parlai.utils.misc import round_sigfigs
  from parlai.agents.transformer.transformer import TransformerGeneratorAgent
  from .modules import EndToEndModel
  from parlai.tasks.wizard_of_wikipedia.agents import TOKEN_KNOWLEDGE
label_desc: Use image node_id 2 for calling the EndToEndAgent obj's underlying member method code with example usage: obj.compute_loss(batch, return_output) and returns: loss,  model_output, loss