Repository: Project-MONAI/MONAI
File: monai/losses/sure_loss.py (Python)

Code before:
        Args:
            operator (function): The operator function that takes in an input
            tensor x and returns an output tensor y. We will use this to compute
            the divergence. More specifically, we will perturb the input x by a
            small amount and compute the divergence between the perturbed output
            and the reference output

            x (torch.Tensor): The input tensor of shape (B, C, H, W) to the
            operator. C=1 or 2: For complex input, the shape is (B, 2, H, W) aka
            C=2 real.  For real input, the shape is (B, 1, H, W) real.

            y_pseudo_gt (torch.Tensor): The pseudo ground truth tensor of shape
            (B, C, H, W) used to compute the L2 loss. C=1 or 2: For complex
            input, the shape is (B, 2, H, W) aka C=2 real.  For real input, the
            shape is (B, 1, H, W) real.

            y_ref (torch.Tensor, optional): The reference output tensor of the
            same shape as y_pseudo_gt

        Returns:
            sure_loss (torch.Tensor): The SURE loss scalar.
        """

        # check inputs shapes
        if x.dim() != 4:
            raise ValueError("Input tensor x should be 4D.")
        if y_pseudo_gt.dim() != 4:
            raise ValueError("Input tensor y_pseudo_gt should be 4D.")
        if y_ref is not None and y_ref.dim() != 4:
            raise ValueError("Input tensor y_ref should be 4D.")
        if x.shape != y_pseudo_gt.shape:
            raise ValueError("Input tensor x and y_pseudo_gt should have the same shape.")
        if y_ref is not None and y_pseudo_gt.shape != y_ref.shape:
            raise ValueError("Input tensor y_pseudo_gt and y_ref should have the same shape.")

        # compute loss
        loss = sure_loss_function(operator, x, y_pseudo_gt, y_ref, self.eps, self.perturb_noise, complex_input)

        return loss


Review comment (suggestion):
```suggestion
            raise ValueError(f"Input tensor x should be 4D, got {x.dim()}.")
```

Code after:
        Args:
            operator (function): The operator function that takes in an input
            tensor x and returns an output tensor y. We will use this to compute
            the divergence. More specifically, we will perturb the input x by a
            small amount and compute the divergence between the perturbed output
            and the reference output

            x (torch.Tensor): The input tensor of shape (B, C, H, W) to the
            operator. C=1 or 2: For complex input, the shape is (B, 2, H, W) aka
            C=2 real.  For real input, the shape is (B, 1, H, W) real.

            y_pseudo_gt (torch.Tensor): The pseudo ground truth tensor of shape
            (B, C, H, W) used to compute the L2 loss. C=1 or 2: For complex
            input, the shape is (B, 2, H, W) aka C=2 real.  For real input, the
            shape is (B, 1, H, W) real.

            y_ref (torch.Tensor, optional): The reference output tensor of the
            same shape as y_pseudo_gt

        Returns:
            sure_loss (torch.Tensor): The SURE loss scalar.
        """

        # check inputs shapes
        if x.dim() != 4:
            raise ValueError(f"Input tensor x should be 4D, got {x.dim()}.")
        if y_pseudo_gt.dim() != 4:
            raise ValueError(f"Input tensor y_pseudo_gt should be 4D, but got {y_pseudo_gt.dim()}.")
        if y_ref is not None and y_ref.dim() != 4:
            raise ValueError(f"Input tensor y_ref should be 4D, but got {y_ref.dim()}.")
        if x.shape != y_pseudo_gt.shape:
            raise ValueError(
                f"Input tensor x and y_pseudo_gt should have the same shape, but got x shape {x.shape}, "
                f"y_pseudo_gt shape {y_pseudo_gt.shape}."
            )
        if y_ref is not None and y_pseudo_gt.shape != y_ref.shape:
            raise ValueError(
                f"Input tensor y_pseudo_gt and y_ref should have the same shape, but got y_pseudo_gt shape {y_pseudo_gt.shape}, "
                f"y_ref shape {y_ref.shape}."
            )

        # compute loss
        loss = sure_loss_function(operator, x, y_pseudo_gt, y_ref, self.eps, self.perturb_noise, complex_input)

        return loss