code: def save(self, model_path, weights_only=False):
      model_state_dict = self.state_dict()
      if weights_only:
          if self.using_tpu:
              xm.save(model_state_dict, model_path)
          else:
              torch.save(model_state_dict, model_path)
          return
      if self.optimizer is not None:
          opt_state_dict = self.optimizer.state_dict()
      else:
          opt_state_dict = None
      if self.scheduler is not None:
          sch_state_dict = self.scheduler.state_dict()
      else:
          sch_state_dict = None
      model_dict = {}
      model_dict["state_dict"] = model_state_dict
      model_dict["optimizer"] = opt_state_dict
      model_dict["scheduler"] = sch_state_dict
      model_dict["epoch"] = self.current_epoch
      model_dict["fp16"] = self.fp16
      if self.using_tpu:
          xm.save(model_dict, model_path)
      else:
          torch.save(model_dict, model_path)
  
imports: import warnings
  import psutil
  import torch
  import torch.nn
  from tqdm import tqdm
  from tez import enums
  from tez.callbacks import CallbackRunner
  from tez.utils import AverageMeter
label_desc: Use image node_id 22 for calling the Model obj's underlying member method code with example usage: obj.save(model_path, weights_only) without return types