class_docstr: Bert-based on PyTorch model for SQuAD-like problem setting:
  It predicts start and end position of answer for given question and context.
  
  [CLS] token is used as no_answer. If model selects [CLS] token as most probable
  answer, it means that there is no answer in given context.
  
  Start and end position of answer are predicted by linear transformation
  of Bert outputs.
  
  Args:
      pretrained_bert: pretrained Bert checkpoint path or key title (e.g. "bert-base-uncased")
      attention_probs_keep_prob: keep_prob for Bert self-attention layers
      hidden_keep_prob: keep_prob for Bert hidden layers
      bert_config_file: path to Bert configuration file, or None, if `pretrained_bert` is a string name
      psg_cls: whether to use a separate linear layer to define if a passage contains the answer to the question
      batch_size: batch size for inference of squad model
docstr: Train model on given batch.
  This method calls train_op using features and labels from y_st and y_end
  
  Args:
      features: batch of InputFeatures instances
      y_st: batch of lists of ground truth answer start positions
      y_end: batch of lists of ground truth answer end positions
  
  Returns:
      dict with loss and learning_rate values
docstr_tok: ["Train","model","on","given","batch",".","This","method","calls","train_op","using","features","and","labels","from","y_st","and","y_end","Args",":","features",":","batch","of","InputFeatures","instances","y_st",":","batch","of","lists","of","ground","truth","answer","start","positions","y_end",":","batch","of","lists","of","ground","truth","answer","end","positions","Returns",":","dict","with","loss","and","learning_rate","values"]
code: def train_on_batch(
      self,
      features: List[List[InputFeatures]],
      y_st: List[List[int]],
      y_end: List[List[int]],
  ) -> Dict:
      """Train model on given batch.
      This method calls train_op using features and labels from y_st and y_end
  
      Args:
          features: batch of InputFeatures instances
          y_st: batch of lists of ground truth answer start positions
          y_end: batch of lists of ground truth answer end positions
  
      Returns:
          dict with loss and learning_rate values
  
      """
      input_ids = [f[0].input_ids for f in features]
      input_masks = [f[0].attention_mask for f in features]
      input_type_ids = [f[0].token_type_ids for f in features]
  
      b_input_ids = torch.cat(input_ids, dim=0).to(self.device)
      b_input_masks = torch.cat(input_masks, dim=0).to(self.device)
      b_input_type_ids = torch.cat(input_type_ids, dim=0).to(
          self.device
      )
  
      y_st = [x[0] for x in y_st]
      y_end = [x[0] for x in y_end]
      b_y_st = torch.from_numpy(np.array(y_st)).to(self.device)
      b_y_end = torch.from_numpy(np.array(y_end)).to(self.device)
  
      input_ = {
          "input_ids": b_input_ids,
          "attention_mask": b_input_masks,
          "token_type_ids": b_input_type_ids,
          "start_positions": b_y_st,
          "end_positions": b_y_end,
          "return_dict": True,
      }
  
      self.optimizer.zero_grad()
      input_ = {
          arg_name: arg_value
          for arg_name, arg_value in input_.items()
          if arg_name in self.accepted_keys
      }
      loss = self.model(**input_).loss
      if self.is_data_parallel:
          loss = loss.mean()
      self._make_step(loss)
  
      return {"loss": loss.item()}
  
imports: from collections import namedtuple
  from logging import getLogger
  from pathlib import Path
  from typing import List, Tuple, Optional, Dict
  import numpy
  import torch
  from transformers import AutoModelForQuestionAnswering, AutoConfig, AutoModel
  from transformers.data.processors.utils import InputFeatures
  from deeppavlov.core.commands.utils import expand_path
  from deeppavlov.core.common.errors import ConfigError
  from deeppavlov.core.common.registry import register
  from deeppavlov.core.models.torch_model import TorchModel
label_desc: Use image node_id 2 for calling the TorchTransformersSquad obj's underlying member method code with example usage: obj.train_on_batch(features, y_st, y_end) and returns: dict