code: def __init__(
      self,
      model,
      dataset,
      device,
      graph_classification=True,
      save_dir=None,
      save_name="model",
      **kwargs,
  ):
      self.model = model
      self.dataset = dataset  # train_mask, eval_mask, test_mask
      self.loader = None
      self.device = device
      self.graph_classification = graph_classification
      self.node_classification = not graph_classification
  
      self.optimizer = None
      self.save = save_dir is not None
      self.save_dir = save_dir
      self.save_name = save_name
      check_dir(self.save_dir)
  
      if self.graph_classification:
          dataloader_params = kwargs.get("dataloader_params")
          self.loader = get_dataloader(dataset, **dataloader_params)
  
imports: import os
  import hydra
  import torch
  import shutil
  import warnings
  from torch.optim import Adam
  from omegaconf import OmegaConf
  from benchmarks.xgraph.utils import check_dir
  from benchmarks.xgraph.gnnNets import get_gnnNets
  from benchmarks.xgraph.dataset import get_dataset, get_dataloader
  import torch.nn.functional
  from torch.optim.lr_scheduler import MultiStepLR
  from utils import fix_random_seed
label_desc: Use image node_id 1 to create a new TrainModel object from inherited base classes: object with example: obj = TrainModel(model, dataset, device, graph_classification, save_dir, save_name)