import torch


def to_tensor(im, dims=3):
  """ Converts a given ndarray image to torch tensor image.

  Args:
    im: ndarray image (height x width x channel x [sample]).
    dims: dimension number of the given image. If dims = 3, the image should
      be in (height x width x channel) format; while if dims = 4, the image
      should be in (height x width x channel x sample) format; default is 3.

  Returns:
    torch tensor in the format (channel x height x width)  or (sample x
      channel x height x width).
  """

  assert (dims == 3 or dims == 4)
  if dims == 3:
    im = im.transpose((2, 0, 1))
  elif dims == 4:
    im = im.transpose((3, 2, 0, 1))
  else:
    raise NotImplementedError
  return torch.from_numpy(im)