    def train_step(self, model: nn.Module, inputs: Any) -> dict:
        """
        The training batch step.
        It contains computing step loss, loss backpropagation, doing step optimization, computing metric.

        Args:
            model (`nn.Module`):
                Pytorch model.
            inputs (`Any`):
                Pytorch tensor or tensor collection.

        Returns:
            (`dict`)
                Step logs which contains the step loss and metric infos.
        """
        step_loss = self.compute_loss(model, inputs)
        step_loss.backward()
        step_loss = step_loss.detach()

        loss_metric, epoch_metric = self.update_metrics(model, inputs, step_loss, training=True)

        self.optimizer.step()
        self.lr_scheduler.step()
        self.optimizer.zero_grad()
        step_logs = {"step_loss": step_loss.item(), "epoch_loss": loss_metric, "epoch_metric": epoch_metric}
        return step_logs
The training batch step.
It contains computing step loss, loss backpropagation, doing step optimization, computing metric.

Args:
model (`nn.Module`):
Pytorch model.
inputs (`Any`):
Pytorch tensor or tensor collection.

Returns:
(`dict`)
Step logs which contains the step loss and metric infos.
towhee/trainer/trainer.py