text: import torch
  import torch.nn as nn
  import numpy as np
  import imageio
  from skimage import img_as_ubyte
  
  def save_images(images, size, image_path):
      return imsave(images, size, image_path)
  
  def imsave(images, size, path):
      image = np.squeeze(merge(images, size))
      return imageio.imwrite(path, img_as_ubyte(image))
  
  def merge(images, size):
      h, w = images.shape[1], images.shape[2]
      if (images.shape[3] in (3,4)):
          c = images.shape[3]
          img = np.zeros((h * size[0], w * size[1], c))
          for idx, image in enumerate(images):
              i = idx % size[1]
              j = idx // size[1]
              img[j * h:j * h + h, i * w:i * w + w, :] = image
          return img
      elif images.shape[3]==1:
          img = np.zeros((h * size[0], w * size[1]))
          for idx, image in enumerate(images):
              i = idx % size[1]
              j = idx // size[1]
              img[j * h:j * h + h, i * w:i * w + w] = image[:,:,0]
          return img
      else:
          raise ValueError('in merge(images,size) images parameter ''must have dimensions: HxW or HxWx3 or HxWx4')
          
  def initialize_weights(net):
      for m in net.modules():
          if isinstance(m, nn.Conv2d):
              m.weight.data.normal_(0, 0.02)
              m.bias.data.zero_()
          elif isinstance(m, nn.ConvTranspose2d):
              m.weight.data.normal_(0, 0.02)
              m.bias.data.zero_()
          elif isinstance(m, nn.Linear):
              m.weight.data.normal_(0, 0.02)
              m.bias.data.zero_()