docstr: Trains the relation extraction BERT model on the given batch.
  Returns:
      dict with loss and learning rate values.
docstr_tok: ["Trains","the","relation","extraction","BERT","model","on","the","given","batch",".","Returns",":","dict","with","loss","and","learning","rate","values","."]
code: def train_on_batch(
      self,
      input_ids: List,
      attention_mask: List,
      entity_pos: List,
      entity_tags: List,
      labels: List,
  ) -> float:
      """
      Trains the relation extraction BERT model on the given batch.
      Returns:
          dict with loss and learning rate values.
      """
  
      _input = {
          "input_ids": torch.LongTensor(input_ids).to(self.device),
          "attention_mask": torch.LongTensor(attention_mask).to(
              self.device
          ),
          "entity_pos": entity_pos,
          "ner_tags": entity_tags,
          "labels": labels,
      }
  
      self.model.train()
      self.model.zero_grad()
      self.optimizer.zero_grad()  # zero the parameter gradients
  
      hidden_states = self.model(**_input)
      loss = hidden_states[0]
      self._make_step(loss)
  
      return loss.item()
  
file_path: DeepPavlov/deeppavlov/models/relation_extraction/relation_extraction_bert.py
imports: from logging import getLogger
  from typing import List, Optional, Union
  import numpy
  import torch
  from deeppavlov.core.common.errors import ConfigError
  from deeppavlov.core.common.registry import register
  from deeppavlov.core.models.torch_model import TorchModel
  from deeppavlov.models.classifiers.re_bert import BertWithAdaThresholdLocContextPooling
label_desc: Use image node_id 2 for calling the REBertModel obj's underlying member method code with example usage: obj.train_on_batch(input_ids, attention_mask, entity_pos, entity_tags, labels) and returns: loss