Repository: Project-MONAI/MONAI
File: monai/engines/evaluator.py (Python)

Code before:
        callback function for the Supervised Evaluation processing logic of 1 iteration in Ignite Engine.
        Return below items in a dictionary:
            - IMAGE: image Tensor data for model input, already moved to device.
            - LABEL: label Tensor data corresponding to the image, already moved to device.
            - PRED: prediction result of model.

        Args:
            engine: `SupervisedEvaluator` to execute operation for an iteration.
            batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data.

        Raises:
            ValueError: When ``batchdata`` is None.

        """
        if batchdata is None:
            raise ValueError("Must provide batch data for current iteration.")
        batch = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs)
        if len(batch) == 2:
            inputs, targets = batch
            args: tuple = ()
            kwargs: dict = {}
        else:
            inputs, targets, args, kwargs = batch
        if self.compile:
            inputs = torch.Tensor(inputs) if isinstance(inputs, MetaTensor) else inputs
            targets = torch.Tensor(targets) if isinstance(targets, MetaTensor) else targets
        # put iteration outputs into engine.state
        engine.state.output = {Keys.IMAGE: inputs, Keys.LABEL: targets}
        # execute forward computation
        with engine.mode(engine.network):
            if engine.amp:
                with torch.cuda.amp.autocast(**engine.amp_kwargs):
                    engine.state.output[Keys.PRED] = engine.inferer(inputs, engine.network, *args, **kwargs)
            else:
                engine.state.output[Keys.PRED] = engine.inferer(inputs, engine.network, *args, **kwargs)
        engine.fire_event(IterationEvents.FORWARD_COMPLETED)
        engine.fire_event(IterationEvents.MODEL_COMPLETED)

        return engine.state.output


class EnsembleEvaluator(Evaluator):
    """
    Ensemble evaluation for multiple models, inherits from evaluator and Workflow.
    It accepts a list of models for inference and outputs a list of predictions for further operations.

    Args:
        device: an object representing the device on which to run.
        val_data_loader: Ignite engine use data_loader to run, must be Iterable, typically be torch.DataLoader.
        epoch_length: number of iterations for one epoch, default to `len(val_data_loader)`.
        networks: networks to evaluate in order in the evaluator, should be regular PyTorch `torch.nn.Module`.

Review comment (suggestion):
Is there a specific reason MetaTensor isn't compatible with compilation? I haven't tried it myself and a bit out of the loop on this. If we cast to regular tensor like this the outputs from networks lose their relationship with the original inputs, eg. filenames and such in the metadata. 

Code after:
        self.compile = compile
        self.inferer = SimpleInferer() if inferer is None else inferer

    def _iteration(self, engine: SupervisedEvaluator, batchdata: dict[str, torch.Tensor]) -> dict:
        """
        callback function for the Supervised Evaluation processing logic of 1 iteration in Ignite Engine.
        Return below items in a dictionary:
            - IMAGE: image Tensor data for model input, already moved to device.
            - LABEL: label Tensor data corresponding to the image, already moved to device.
            - PRED: prediction result of model.

        Args:
            engine: `SupervisedEvaluator` to execute operation for an iteration.
            batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data.

        Raises:
            ValueError: When ``batchdata`` is None.

        """
        if batchdata is None:
            raise ValueError("Must provide batch data for current iteration.")
        batch = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs)
        if len(batch) == 2:
            inputs, targets = batch
            args: tuple = ()
            kwargs: dict = {}
        else:
            inputs, targets, args, kwargs = batch
        # FIXME: workaround for https://github.com/pytorch/pytorch/issues/117026
        if self.compile:
            inputs_meta, targets_meta, inputs_applied_operations, targets_applied_operations = None, None, None, None
            if isinstance(inputs, MetaTensor):
                warnings.warn(
                    "Will convert to PyTorch Tensor if using compile, and casting back to MetaTensor after the forward pass."
                )
                inputs, inputs_meta, inputs_applied_operations = (
                    inputs.as_tensor(),
                    inputs.meta,
                    inputs.applied_operations,
                )
            if isinstance(targets, MetaTensor):
                targets, targets_meta, targets_applied_operations = (
                    targets.as_tensor(),
                    targets.meta,
                    targets.applied_operations,
                )

        # put iteration outputs into engine.state
        engine.state.output = {Keys.IMAGE: inputs, Keys.LABEL: targets}
        # execute forward computation
        with engine.mode(engine.network):