def train_on_batch(self, input_ids: Union[List[List[int]], np.ndarray],
                       input_masks: Union[List[List[int]], np.ndarray],
                       y_masks: Union[List[List[int]], np.ndarray],
                       y_heads: List[List[int]], y_dep: List[List[int]]) -> Dict:
        """

        Args:
            input_ids: indices of the subwords
            input_masks: mask that determines where to attend and where not to
            y_masks: mask which determines the first subword units in the the word
            y_heads: for each token - id fo token which is the head in syntax tree for the token
            y_dep: syntax dependencies for each tokens
        """
        self.optimizer.zero_grad()
        loss = self.model(input_ids, input_masks, y_masks, y_heads, y_dep)
        loss.backward()
        # Clip the norm of the gradients to 1.0.
        # This is to help prevent the "exploding gradients" problem.
        if self.clip_norm:
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip_norm)

        self.optimizer.step()

        return {'loss': loss.item()}


        Args:
            input_ids: indices of the subwords
            input_masks: mask that determines where to attend and where not to
            y_masks: mask which determines the first subword units in the the word
            y_heads: for each token - id fo token which is the head in syntax tree for the token
            y_dep: syntax dependencies for each tokens
        
deeppavlov/models/torch_bert/torch_transformers_syntax_parser.py