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

        if self.n_classes > 1 and not self.is_binary:
            _input["labels"] = torch.from_numpy(np.array(y)).to(self.device)

        # regression
        else:
            _input["labels"] = torch.from_numpy(np.array(y, dtype=np.float32)).unsqueeze(1).to(self.device)

        self.optimizer.zero_grad()

        tokenized = {key: value for (key, value) in _input.items()
                     if key in self.accepted_keys}

        loss = self.model(**tokenized).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 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
        
deeppavlov/models/torch_bert/torch_transformers_classifier.py