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.
docstr: Sets up the devices to perform distributed training.
docstr_tok: ["Sets","up","the","devices","to","perform","distributed","training","."]
code: def _setup_devices(self):
      """Sets up the devices to perform distributed training."""
  
      if (
          dist.is_available()
          and dist.is_initialized()
          and self.local_rank == -1
      ):
          logger.warning(
              "torch.distributed process group is initialized, but local_rank == -1. "
          )
      if self.training_config.no_cuda:
          self._n_gpus = 0
          device = "cpu"
  
      else:
          torch.cuda.set_device(self.local_rank)
          device = torch.device("cuda", self.local_rank)
  
          if not dist.is_initialized():
              dist.init_process_group(
                  backend=self.dist_backend,
                  init_method="env://",
                  world_size=self.world_size,
                  rank=self.rank,
              )
  
      return device
  
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 3 for calling the BaseTrainer obj's underlying member method code with example usage: obj._setup_devices() and returns: device