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

Code before:
            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
        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}, 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}, y_ref shape {y_ref.shape}.")
```

Code after:
            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