def train_on_batch(self, *args):
        """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)
        Returns:
            dict with loss for each task
        """
        log.debug(f'Training for {args}')
        error_msg = f'Len of arguments {len(args)} is WRONG. ' \
                    f'Correct is {2 * self.n_tasks} as n_tasks is {self.n_tasks}'
        if len(args) != 2 * self.n_tasks:
            raise Exception(error_msg)
        ids_to_iterate = [k for k in range(self.n_tasks) if len(args[k]) > 0]
        if len(ids_to_iterate) == 0:
            raise Exception(f'No examples given! Given args {args}')
        elif len(ids_to_iterate) > 1:
            raise Exception('Samples from more than 1 task in train_on_batch')
        task_id = ids_to_iterate[0]
        _input, batch_size = self._make_input(task_features=args[task_id], task_id=task_id,
                                              labels=args[task_id + self.n_tasks])
        if _input == {}:
            raise Exception('Empty input!')

        if self.prev_id is None:
            self.prev_id = task_id
        elif self.prev_id != task_id and not self.printed:
            log.info('Seen samples from different tasks')
            self.printed = True
        if 'token_type_ids' not in _input:
            _input['token_type_ids'] = None
        loss, logits = self.model(task_id=task_id, **_input)
        if self.is_data_parallel:
            loss = loss.mean()
        loss = loss / self.gradient_accumulation_steps
        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)

        if (self.steps_taken + 1) % self.gradient_accumulation_steps == 0 or (
                self.steps_per_epoch is not None and (self.steps_taken + 1) % self.steps_per_epoch == 0):
            self.optimizer.step()
            self.optimizer.zero_grad()
        self.train_losses[task_id] = loss.item()
        self.steps_taken += 1
        log.debug(f'train {task_id} {logits}')
        return {"losses": self.train_losses}
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)
        Returns:
            dict with loss for each task
        
deeppavlov/models/torch_bert/multitask_transformer.py