code: def _init_trainer(
      self, train_dataset, valid_dataset, config, **kwargs
  ):
      self.config = config
      self.train_dataset = train_dataset
      self.valid_dataset = valid_dataset
      self._init_driver()
  
      # parse args
      self.train_loader = kwargs.get("train_loader", None)
      self.valid_loader = kwargs.get("valid_loader", None)
      self.train_sampler = kwargs.get("train_sampler", None)
      self.valid_sampler = kwargs.get("valid_sampler", None)
      self.train_collate_fn = kwargs.get("train_collate_fn", None)
      self.valid_collate_fn = kwargs.get("valid_collate_fn", None)
  
      if self.config.num_jobs == -1:
          self.config.num_jobs = multiprocessing.cpu_count()
          if self.config.num_jobs > 4:
              self.config.num_jobs -= 2
  
      if self.train_loader is None:
          self.train_loader = DataLoader(
              self.train_dataset,
              batch_size=self.config.training_batch_size,
              num_workers=self.config.num_jobs,
              sampler=self.train_sampler,
              shuffle=self.config.train_shuffle,
              collate_fn=self.train_collate_fn,
              drop_last=self.config.train_drop_last,
              pin_memory=self.config.pin_memory,
          )
  
      if self.valid_loader is None:
          if self.valid_dataset is not None:
              self.valid_loader = DataLoader(
                  self.valid_dataset,
                  batch_size=self.config.validation_batch_size,
                  num_workers=self.config.num_jobs,
                  sampler=self.valid_sampler,
                  shuffle=self.config.valid_shuffle,
                  collate_fn=self.valid_collate_fn,
                  drop_last=self.config.valid_drop_last,
                  pin_memory=self.config.pin_memory,
              )
  
      self.optimizer, self.scheduler = self.model.optimizer_scheduler()
  
      if self.optimizer is None:
          raise Exception("No optimizer found")
  
      if self.valid_loader is not None and self.scheduler is not None:
          (
              self.model,
              self.optimizer,
              self.train_loader,
              self.valid_loader,
              self.scheduler,
          ) = self._driver.prepare(
              self.model,
              self.optimizer,
              self.train_loader,
              self.valid_loader,
              self.scheduler,
          )
      elif self.valid_loader is not None and self.scheduler is None:
          (
              self.model,
              self.optimizer,
              self.train_loader,
              self.valid_loader,
          ) = self._driver.prepare(
              self.model,
              self.optimizer,
              self.train_loader,
              self.valid_loader,
          )
      elif self.valid_loader is None and self.scheduler is not None:
          (
              self.model,
              self.optimizer,
              self.train_loader,
              self.scheduler,
          ) = self._driver.prepare(
              self.model,
              self.optimizer,
              self.train_loader,
              self.scheduler,
          )
      else:
          (
              self.model,
              self.optimizer,
              self.train_loader,
          ) = self._driver.prepare(
              self.model, self.optimizer, self.train_loader
          )
  
      self.num_train_steps = int(
          len(self.train_loader) * self.config.epochs
      )
      if self.valid_dataset:
          self.num_valid_steps = len(self.valid_loader)
      else:
          self.num_valid_steps = None
  
      self._progress = Progress(
          num_train_steps=self.num_train_steps,
          num_valid_steps=self.num_valid_steps,
      )
  
      if "callbacks" in kwargs:
          self.callbacks = [self._progress] + kwargs["callbacks"]
      else:
          self.callbacks = [self._progress]
  
      self._callback_runner = CallbackRunner(self.callbacks, self)
      self.train_state = enums.TrainingState.TRAIN_START
  
imports: import multiprocessing
  import time
  import warnings
  from dataclasses import dataclass
  from typing import Optional
  import torch
  from accelerate import Accelerator
  from torch.utils.data import DataLoader
  from tez import enums
  from tez.callbacks import CallbackRunner, Progress
  from tez.logger import logger
  from tez.utils import AverageMeter
  from .config import TezConfig
label_desc: Use image node_id 2 for calling the Tez obj's underlying member method code with example usage: obj._init_trainer(train_dataset, valid_dataset, config) without return types