def METHOD_NAME(self, features_li: List[List[InputFeatures]], y: Union[List[int], List[List[int]]]) -> Dict:
    """Train the model on the given batch.
    Args:
        features_li: list with the single element containing the batch of InputFeatures
        y: batch of labels (class id or one-hot encoding)
    Returns:
        dict with loss and learning rate values
    """
    features = features_li[0]
    input_ids = [f.input_ids for f in features]
    input_masks = [f.attention_mask 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_labels = torch.from_numpy(np.array(y)).to(self.device)
    self.optimizer.zero_grad()
    loss, logits = self.model(b_input_ids, token_type_ids=None, attention_mask=b_input_masks,
                              labels=b_labels, return_dict=False)
    self._make_step(loss)
    return {'loss': loss.item()}