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()}
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

        
deeppavlov/models/torch_bert/torch_transformers_squad.py