import torch


def to_4D_image(image: torch.Tensor) -> torch.Tensor:
  """Convert `image` to 4D tensors (b, c, h, w)

  Args:
    image (torch.Tensor): 2/3/4D image tensor
      2D: (h, w)
      3D: (c, h, w)
      4D: (b, c, h, w)
  
  Returns:
    torch.Tensor: 4D image tensor
  """
  ndims = len(image.shape)
  assert ndims in [2, 3, 4], \
    f"only supports 2/3/4D images while {ndims}-D are given."
  if ndims == 2:
    return image[None, None, :, :] # b, c
  elif ndims == 3:
    return image[None, :, :, :] # b
  else:
    return image