code: def compute_qa_loss(logits, start_positions, end_positions):
      # Do we want to keep them as 1 tensor, or multiple?
      # bs x 2 x seq_len x 1
  
      start_logits, end_logits = logits[:, 0], logits[:, 1]
      # Taken from: RobertaForQuestionAnswering
      # If we are on multi-GPU, split add a dimension
      if len(start_positions.size()) > 1:
          start_positions = start_positions.squeeze(-1)
      if len(end_positions.size()) > 1:
          end_positions = end_positions.squeeze(-1)
      # sometimes the start/end positions are outside our model inputs, we ignore these terms
      ignored_index = start_logits.size(1)
      start_positions.clamp_(0, ignored_index)
      end_positions.clamp_(0, ignored_index)
  
      loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index)
      start_loss = loss_fct(start_logits, start_positions)
      end_loss = loss_fct(end_logits, end_positions)
      total_loss = (start_loss + end_loss) / 2
      return total_loss
  
imports: import abc
  import torch
  import torch.nn
  from typing import Callable
  import jiant.proj.main.modeling.heads
  from jiant.proj.main.components.outputs import LogitsAndLossOutput
  from jiant.proj.main.components.outputs import LogitsOutput
  from jiant.utils.python.datastructures import take_one
  from jiant.tasks.core import TaskTypes
label_desc: Use image node_id 2 for calling a global function with example usage: compute_qa_loss(logits, start_positions, end_positions) and returns: total_loss