code: def save_image(data, epoch, image_size, batch_size, output_dir, padding=2):
      
      data = data.asnumpy().transpose((0, 2, 3, 1))
      datanp = np.clip(
          (data - np.min(data))*(255.0/(np.max(data) - np.min(data))), 0, 255).astype(np.uint8)
      x_dim = min(8, batch_size)
      y_dim = int(math.ceil(float(batch_size) / x_dim))
      height, width = int(image_size + padding), int(image_size + padding)
      grid = np.zeros((height * y_dim + 1 + padding // 2, width *
                       x_dim + 1 + padding // 2, 3), dtype=np.uint8)
      k = 0
      for y in range(y_dim):
          for x in range(x_dim):
              if k >= batch_size:
                  break
              start_y = y * height + 1 + padding // 2
              end_y = start_y + height - padding
              start_x = x * width + 1 + padding // 2
              end_x = start_x + width - padding
              np.copyto(grid[start_y:end_y, start_x:end_x, :], datanp[k])
              k += 1
      imageio.imwrite(
          '{}/fake_samples_epoch_{}.png'.format(output_dir, epoch), grid)
docs: save image