class_docstr: Base class to perform model training.
  
  Args:
      model (BaseAE): A instance of :class:`~pythae.models.BaseAE` to train
  
      train_dataset (BaseDataset): The training dataset of type
          :class:`~pythae.data.dataset.BaseDataset`
  
      eval_dataset (BaseDataset): The evaluation dataset of type
          :class:`~pythae.data.dataset.BaseDataset`
  
      training_config (BaseTrainerConfig): The training arguments summarizing the main
          parameters used for training. If None, a basic training instance of
          :class:`BaseTrainerConfig` is used. Default: None.
  
      callbacks (List[~pythae.trainers.training_callbacks.TrainingCallbacks]):
          A list of callbacks to use during training.
code: def __init__(
      self,
      model: BaseAE,
      train_dataset: Union[BaseDataset, DataLoader],
      eval_dataset: Optional[Union[BaseDataset, DataLoader]] = None,
      training_config: Optional[BaseTrainerConfig] = None,
      callbacks: List[TrainingCallback] = None,
  ):
      if training_config is None:
          training_config = BaseTrainerConfig()
  
      if training_config.output_dir is None:
          output_dir = "dummy_output_dir"
          training_config.output_dir = output_dir
  
      self.training_config = training_config
      self.model_config = model.model_config
      self.model_name = model.model_name
  
      # for distributed training
      self.world_size = self.training_config.world_size
      self.local_rank = self.training_config.local_rank
      self.rank = self.training_config.rank
      self.dist_backend = self.training_config.dist_backend
  
      if self.world_size > 1:
          self.distributed = True
      else:
          self.distributed = False
  
      if self.distributed:
          device = self._setup_devices()
  
      else:
          device = (
              "cuda"
              if torch.cuda.is_available()
              and not self.training_config.no_cuda
              else "cpu"
          )
  
      self.amp_context = (
          torch.autocast("cuda")
          if self.training_config.amp
          else contextlib.nullcontext()
      )
  
      if (
          hasattr(model.model_config, "reconstruction_loss")
          and model.model_config.reconstruction_loss == "bce"
      ):
          self.amp_context = contextlib.nullcontext()
  
      self.device = device
  
      # place model on device
      model = model.to(device)
      model.device = device
  
      if self.distributed:
          model = DDP(model, device_ids=[self.local_rank])
  
      self.train_dataset = train_dataset
      self.eval_dataset = eval_dataset
  
      # Define the loaders
      if isinstance(train_dataset, DataLoader):
          train_loader = train_dataset
          logger.warn(
              "Using the provided train dataloader! Carefull this may overwrite some "
              "parameters provided in your training config."
          )
      else:
          train_loader = self.get_train_dataloader(train_dataset)
  
      if eval_dataset is not None:
          if isinstance(eval_dataset, DataLoader):
              eval_loader = eval_dataset
              logger.warn(
                  "Using the provided eval dataloader! Carefull this may overwrite some "
                  "parameters provided in your training config."
              )
          else:
              eval_loader = self.get_eval_dataloader(eval_dataset)
      else:
          logger.info(
              "! No eval dataset provided ! -> keeping best model on train.\n"
          )
          self.training_config.keep_best_on_train = True
          eval_loader = None
  
      self.train_loader = train_loader
      self.eval_loader = eval_loader
      self.callbacks = callbacks
  
      # run sanity check on the model
      self._run_model_sanity_check(model, train_loader)
  
      if self.is_main_process:
          logger.info(
              "Model passed sanity check !\n" "Ready for training.\n"
          )
  
      self.model = model
  
imports: import contextlib
  import datetime
  import logging
  import os
  from copy import deepcopy
  from typing import Any, Dict, List, Optional, Union
  import torch
  import torch.distributed
  import torch.optim
  import torch.optim.lr_scheduler
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data import DataLoader, Dataset
  from torch.utils.data.distributed import DistributedSampler
  from ...customexception import ModelError
  from ...data.datasets import BaseDataset, collate_dataset_output
  from ...models import BaseAE
  from ..trainer_utils import set_seed
  from ..training_callbacks import CallbackHandler, MetricConsolePrinterCallback, ProgressBarCallback, TrainingCallback
  from .base_training_config import BaseTrainerConfig
label_desc: Use image node_id 1 to create a new BaseTrainer object  with example: obj = BaseTrainer(model, train_dataset, eval_dataset, training_config, callbacks)