def METHOD_NAME(self, batch: Any, batch_idx: int) -> Any:
    """For the test 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.test_step`."""
    batch = (batch[DataKeys.INPUT], batch[DataKeys.TARGET])
    return super().METHOD_NAME(batch, batch_idx)