def METHOD_NAME(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().METHOD_NAME(batch, batch_idx, dataloader_idx=dataloader_idx)