def METHOD_NAME(self, features: Dict[str, torch.tensor], y: Union[List[int], List[List[int]]]) -> Dict:
    """Train model on given batch.
    This method calls train_op using features and y (labels).
    Args:
        features: batch of InputFeatures
        y: batch of labels (class id or one-hot encoding)
    Returns:
        dict with loss and learning_rate values
    """
    _input = {key: value.to(self.device) for key, value in features.items()}
    _input["labels"] = torch.tensor(y).long().to(self.device)
    self.optimizer.zero_grad()
    tokenized = {key: value for (key, value) in _input.items() if key in self.model.forward.__code__.co_varnames}
    loss = self.model(**tokenized).loss
    self._make_step(loss)
    return {'loss': loss.item()}