code: def load(self, model_path, weights_only=False, device="cuda"):
      if device == "tpu":
          if XLA_AVAILABLE is False:
              raise RuntimeError("XLA is not available")
          else:
              self.using_tpu = True
              device = xm.xla_device()
      self.device = device
      if next(self.parameters()).device != self.device:
          self.to(self.device)
      model_dict = torch.load(
          model_path, map_location=torch.device(device)
      )
      if weights_only:
          self.load_state_dict(model_dict)
      else:
          self.load_state_dict(model_dict["state_dict"])
  
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 23 for calling the Model obj's underlying member method code with example usage: obj.load(model_path, weights_only, device) without return types