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,
      )
  
file_path: atspy/atspy/nbeats.py
imports: import torch
  from torch import optim
  from torch.nn import functional
  from nbeats_pytorch.model import NBeatsNet
  import os
  import numpy
  import matplotlib.pyplot
label_desc: Use image node_id 5 for calling a global function with example usage: save(model, optimiser, grad_step) without return types