code: def save(model, optimiser, grad_step):
      torch.save(
          {
              "grad_step": grad_step,
              "model_state_dict": model.state_dict(),
              "optimizer_state_dict": optimiser.state_dict(),
          },
          CHECKPOINT_NAME,
      )
  
imports: import numpy
  import pandas
  from numpy import matrix
  from scipy import linalg
  from matplotlib.pylab import rcParams
label_desc: Use image node_id 5 for calling a global function with example usage: save(model, optimiser, grad_step) without return types