import torch


def complex_abs(tensor):
  """Compute absolute value of complex image tensor

  Parameters
  ----------
  tensor : torch.Tensor
    Tensor of shape (batch, 2, height, width)

  Returns
  -------
  Tensor with magnitude image of shape (batch, 1, height, width)
  """
  tensor = (tensor[:, 0] ** 2 + tensor[:, 1] ** 2) ** 0.5
  return torch.unsqueeze(tensor, dim=1)