code: def plot_reconstructions(
      dataset_obj,
      trained_net,
      input_images,
      image_id_to_plot=0,
      epoch=0,
      batch_id=0,
      save_dir="results",
  ):
      z_sample, _, _ = trained_net.encode(input_images)
      # print(z_sample[image_id_to_plot])
      recons = trained_net.decode(z_sample)
  
      input_images_numpy = convert_and_reshape(
          input_images, dataset_obj
      )
      recons_numpy = convert_and_reshape(recons, dataset_obj)
  
      f, axarr = plt.subplots(1, 2)
      if input_images_numpy.shape[2] == 1:
          axarr[0].imshow(
              input_images_numpy[image_id_to_plot, :, :, 0], cmap="gray"
          )
          axarr[1].imshow(
              recons_numpy[image_id_to_plot, :, :, 0], cmap="gray"
          )
      else:
          axarr[0].imshow(
              input_images_numpy[image_id_to_plot] * 0.5 + 0.5
          )
          axarr[1].imshow(recons_numpy[image_id_to_plot] * 0.5 + 0.5)
  
      if not os.path.isdir(save_dir):
          os.makedirs(save_dir)
  
      f.savefig(
          os.path.join(
              save_dir,
              "recons_epoch_{}_batch_id_{}.png".format(epoch, batch_id),
          )
      )
      plt.close(fig=f)
  
imports: import matplotlib
  import matplotlib.pyplot
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.autograd import Variable
  from torch.nn import Linear
  import os
label_desc: Use image node_id 2 for calling a global function with example usage: plot_reconstructions(dataset_obj, trained_net, input_images, image_id_to_plot, epoch, batch_id, save_dir) without return types