code: def _prepare_image(img, nrow=8, padding=2, square_image=False):
      
      if isinstance(img, np.ndarray):
          img = nd.array(img, dtype=img.dtype, ctx=current_context())
      if not isinstance(img, NDArray):
          raise TypeError('expected MXNet NDArray or numpy.ndarray, '
                          'while received type {}'.format(str(type(img))))
      assert img.ndim == 2 or img.ndim == 3 or img.ndim == 4
  
      if img.dtype == np.uint8:
          return make_image_grid(
              img, nrow=nrow, padding=padding, square_image=square_image).transpose((1, 2, 0))
      elif img.dtype == np.float32 or img.dtype == np.float64:
          min_val = img.min().asscalar()
          max_val = img.max().asscalar()
          if min_val < 0.0:
              raise ValueError('expected non-negative min value from img, '
                               'while received {}'.format(min_val))
          if max_val > 1.0:
              raise ValueError('expected max value from img not greater than 1, '
                               'while received {}'.format(max_val))
          img = make_image_grid(img, nrow=nrow, padding=padding, square_image=square_image) * 255.0
          return img.astype(np.uint8).transpose((1, 2, 0))
      else:
          raise ValueError('expected input image dtype is one of uint8, float32, '
                           'and float64, received dtype {}'.format(str(img.dtype)))
docs: Given an image of format HW, CHW, or NCHW, returns a image of format HWC.
      If the input is a batch of images, a grid of images is made by stitching them together.
      If data type is float, values must be in the range [0, 1], and then they are rescaled to
      range [0, 255]. If data type is 'uint8`, values are unchanged.