import torch


def vis_tensor(tensor, norm=True, dim=None, scale=1):
  """ Returns a processed tensor for visualization purposes.

  Args:
    tensor: image tensor; if it is in the format (batch x channel x height x
      width), use dim=None (default); otherwise, use dim to determine which
      axis is used for dimension extension to be in the (batch x channel x
      height x width) format.
    norm: boolean to apply min-max normalization; default is True.
    dim: a dimension of size one inserted at the specified position; this
      dimension extension is an optional; default is without extension.
    scale: gain scale for visualization purposes; default is 1.

  Returns:
    enhanced tensor for visualization purposes.
  """

  if dim is not None:
    tensor = torch.unsqueeze(tensor, dim=dim)
  if norm:
    tensor = (tensor - torch.min(tensor)) / (
        torch.max(tensor) - torch.min(tensor))
  tensor = tensor * scale
  return tensor