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