code: def generate_plots(xx, yy, zz, model_name, results_loc):
      zz = np.log(zz)
  
      plt.figure(figsize=(10, 10))
      plt.contour(xx, yy, zz)
      plt.savefig(
          f"{results_loc}/{model_name}_log_contour.png", dpi=100
      )
      plt.close()
  
      ## 3D plot
      fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
      ax.set_axis_off()
      surf = ax.plot_surface(
          xx, yy, zz, cmap=cm.coolwarm, linewidth=0, antialiased=False
      )
      ax.set_xlim(-1, 1)
      ax.set_ylim(-1, 1)
  
      plt.savefig(
          f"{results_loc}/{model_name}_log_surface.png",
          dpi=100,
          format="png",
          bbox_inches="tight",
      )
      plt.close()
  
      fig = plt.figure(figsize=(10, 10))
      ax = Axes3D(fig)
      ax.set_axis_off()
  
      def init():
          ax.plot_surface(
              xx,
              yy,
              zz,
              cmap=cm.coolwarm,
              linewidth=0,
              antialiased=False,
          )
          ax.set_xlim(-1, 1)
          ax.set_ylim(-1, 1)
          return (fig,)
  
      def animate(i):
          ax.view_init(elev=(15 * (i // 15) + i % 15) + 0.0, azim=i)
          ax.set_xlim(-1, 1)
          ax.set_ylim(-1, 1)
          return (fig,)
  
      anim = animation.FuncAnimation(
          fig,
          animate,
          init_func=init,
          frames=100,
          interval=20,
          blit=True,
      )
  
      anim.save(
          f"{results_loc}/{model_name}_log_surface.gif",
          fps=15,
          writer="imagemagick",
      )
  
file_path: ml-cvnets/loss_landscape/landscape_utils.py
imports: import copy
  from typing import Dict, Optional
  import matplotlib.pyplot
  import numpy
  import torch
  from matplotlib import animation, cm
  from mpl_toolkits.mplot3d import Axes3D
  from utils.common_utils import unwrap_model_fn
label_desc: Use image node_id 5 for calling a global function with example usage: generate_plots(xx, yy, zz, model_name, results_loc) and returns: fig, fig