Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/metrics/regression/explained_variance.py (Python)

Code before:
        >>> explained_variance(preds, target)
        tensor(0.9572)

        >>> target = torch.tensor([[0.5, 1], [-1, 1], [7, -6]])
        >>> preds = torch.tensor([[0, 2], [-1, 2], [8, -5]])
        >>> explained_variance = ExplainedVariance(multioutput='raw_values')
        >>> explained_variance(preds, target)
        tensor([0.9677, 1.0000])
    """

    def __init__(
        self,
        multioutput: str = 'uniform_average',
        compute_on_step: bool = True,
        ddp_sync_on_step: bool = False,
        process_group: Optional[Any] = None,
    ):
        super().__init__(
            compute_on_step=compute_on_step,
            ddp_sync_on_step=ddp_sync_on_step,
            process_group=process_group,
        )
        allowed_multioutput = ('raw_values', 'uniform_average', 'variance_weighted')
        assert multioutput in allowed_multioutput, \
            'Invalid input to argument `multioutput`. Choose one of the following:' \
            f'{allowed_multioutput}'
        self.multioutput = multioutput
        self.add_state("y", default=[], dist_reduce_fx=None)
        self.add_state("y_pred", default=[], dist_reduce_fx=None)

        rank_zero_warn('Metric `ExplainedVariance` will save all targets and'
                       ' predictions in buffer. For large datasets this may lead'
                       ' to large memory footprint.')

    def update(self, preds: torch.Tensor, target: torch.Tensor):
        """
        Update state with predictions and targets.

        Args:
            preds: Predictions from model
            target: Ground truth values
        """
        assert preds.shape == target.shape, \
            'Predictions and targets are expected to have the same shape'
        self.y.append(target)
        self.y_pred.append(preds)

    def compute(self):
        """
        Computes explained variance over state.
        """

Review comment (suggestion):
```suggestion
        if multioutput not in allowed_multioutput:
            raise ValueError(
                f'Invalid input to argument `multioutput`. Choose one of the following: {allowed_multioutput}'
            )
```

Code after:
        >>> explained_variance(preds, target)
        tensor(0.9572)

        >>> target = torch.tensor([[0.5, 1], [-1, 1], [7, -6]])
        >>> preds = torch.tensor([[0, 2], [-1, 2], [8, -5]])
        >>> explained_variance = ExplainedVariance(multioutput='raw_values')
        >>> explained_variance(preds, target)
        tensor([0.9677, 1.0000])
    """

    def __init__(
        self,
        multioutput: str = 'uniform_average',
        compute_on_step: bool = True,
        ddp_sync_on_step: bool = False,
        process_group: Optional[Any] = None,
    ):
        super().__init__(
            compute_on_step=compute_on_step,
            ddp_sync_on_step=ddp_sync_on_step,
            process_group=process_group,
        )
        allowed_multioutput = ('raw_values', 'uniform_average', 'variance_weighted')
        if multioutput not in allowed_multioutput:
            raise ValueError(
                f'Invalid input to argument `multioutput`. Choose one of the following: {allowed_multioutput}'
            )
        self.multioutput = multioutput
        self.add_state("y", default=[], dist_reduce_fx=None)
        self.add_state("y_pred", default=[], dist_reduce_fx=None)

        rank_zero_warn('Metric `ExplainedVariance` will save all targets and'
                       ' predictions in buffer. For large datasets this may lead'
                       ' to large memory footprint.')

    def update(self, preds: torch.Tensor, target: torch.Tensor):
        """
        Update state with predictions and targets.

        Args:
            preds: Predictions from model
            target: Ground truth values
        """
        self._check_same_shape(preds, target)
        self.y.append(target)
        self.y_pred.append(preds)

    def compute(self):
        """
        Computes explained variance over state.
        """