def training_step(self, batch: Any, batch_idx: int) -> Any:
        """For the training step, we just extract the :attr:`~flash.core.data.io.input.DataKeys.INPUT` and
        :attr:`~flash.core.data.io.input.DataKeys.TARGET` keys from the input and forward them to the
        :meth:`~flash.core.model.Task.training_step`."""
        batch = (batch[DataKeys.INPUT], batch[DataKeys.TARGET])
        return super().training_step(batch, batch_idx)
For the training step, we just extract the :attr:`~flash.core.data.io.input.DataKeys.INPUT` and
        :attr:`~flash.core.data.io.input.DataKeys.TARGET` keys from the input and forward them to the
        :meth:`~flash.core.model.Task.training_step`.
src/flash/template/classification/model.py