def _tensor_to_list_of_channel_tensors(img):
  """Converts a tensor with dimensions of HWC or BHWC to a list of channels."""
  if len(img.shape) == 3:  # HWC.
    img_list = [img[:, :, c] for c in range(img.shape[-1])]
  elif len(img.shape) == 4:  # BHWC.
    img_list = [img[:, :, :, c] for c in range(img.shape[-1])]
  else:
    raise ValueError(f'Input must be HWC or BHWC and not of shape {img.shape}.')
  return img_list