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

Code before:
import torch.nn as nn
from torch.nn.modules.loss import _Loss


def complex_diff_abs_loss(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    First compute the difference in the complex domain,
    then get the absolute value and take the mse

    Args:
        x, y - B, 2, H, W real valued tensors representing complex numbers
                or  B,1,H,W complex valued tensors
    Returns:
        l2_loss - scalar
    """
    if not x.is_complex():
        x_ = torch.view_as_complex(x.permute(0, 2, 3, 1).contiguous())
    else:
        x_ = x
    if not y.is_complex():
        y_ = torch.view_as_complex(y.permute(0, 2, 3, 1).contiguous())
    else:
        y_ = y
    diff = x_ - y_
    l2_loss = nn.functional.mse_loss(torch.abs(diff), torch.zeros_like(torch.abs(diff)), reduction="mean")
    return l2_loss


def sure_loss_function(
    operator: Callable,
    x: torch.Tensor,
    y_pseudo_gt: torch.Tensor,
    y_ref: Optional[torch.Tensor] = None,
    eps: Optional[float] = -1.0,
    perturb_noise: Optional[torch.Tensor] = None,
    complex_input: Optional[bool] = False,
) -> torch.Tensor:
    """
    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.  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.  For complex input, the shape is

Review comment (refactor):
```suggestion
    if not x.is_complex():
        x = torch.view_as_complex(x.permute(0, 2, 3, 1).contiguous())
    if not y.is_complex():
        y = torch.view_as_complex(y.permute(0, 2, 3, 1).contiguous())
        
    diff = torch.abs(x - y)
    return nn.functional.mse_loss(diff, torch.zeros_like(diff), reduction="mean")
```

I think it's fine to simplify a bit like this.

Code after:
import torch.nn as nn
from torch.nn.modules.loss import _Loss


def complex_diff_abs_loss(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    First compute the difference in the complex domain,
    then get the absolute value and take the mse

    Args:
        x, y - B, 2, H, W real valued tensors representing complex numbers
                or  B,1,H,W complex valued tensors
    Returns:
        l2_loss - scalar
    """
    if not x.is_complex():
        x = torch.view_as_complex(x.permute(0, 2, 3, 1).contiguous())
    if not y.is_complex():
        y = torch.view_as_complex(y.permute(0, 2, 3, 1).contiguous())

    diff = torch.abs(x - y)
    return nn.functional.mse_loss(diff, torch.zeros_like(diff), reduction="mean")


def sure_loss_function(
    operator: Callable,
    x: torch.Tensor,
    y_pseudo_gt: torch.Tensor,
    y_ref: Optional[torch.Tensor] = None,
    eps: Optional[float] = -1.0,
    perturb_noise: Optional[torch.Tensor] = None,
    complex_input: Optional[bool] = False,
) -> torch.Tensor:
    """
    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.  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.  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 shape