class_docstr: Fits a second VAE in the Autoencoder's latent space.
  
  Args:
      model (VAE): The VAE model to sample from
      sampler_config (TwoStageVAESamplerConfig): A TwoStageVAESamplerConfig instance containing
          the main parameters of the sampler. If None, a pre-defined configuration is used.
          Default: None
  
  .. note::
  
      The method :class:`~pythae.samplers.TwoStageVAESampler.fit` must be called to fit the
      sampler before sampling.
class_docstr_tok: ["Fits","a","second","VAE","in","the","Autoencoder","'s","latent","space",".","Args",":","model","(","VAE",")",":","The","VAE","model","to","sample","from","sampler_config","(","TwoStageVAESamplerConfig",")",":","A","TwoStageVAESamplerConfig","instance","containing","the","main","parameters","of","the","sampler",".","If","None",",","a","pre-defined","configuration","is","used",".","Default",":","None","..","note",":",":","The","method",":","class",":","`","~pythae.samplers.TwoStageVAESampler.fit","`","must","be","called","to","fit","the","sampler","before","sampling","."]
docstr: Method to fit the sampler from the training data
  
  Args:
      train_data (Union[torch.Tensor, np.ndarray, Dataset]): The train data needed to
          retrieve the training embeddings and fit the second VAE in the latent space.
      eval_data (Union[torch.Tensor, np.ndarray, Dataset]): The train data needed to retrieve
          the evaluation embeddings and fit the second VAE in the latent space.
      training_config (BaseTrainerConfig): the training config to use to fit the second VAE.
      batch_size (int): The batch size to use to retrieve the embeddings. Default: 64.
docstr_tok: ["Method","to","fit","the","sampler","from","the","training","data","Args",":","train_data","(","Union","[","torch.Tensor",",","np.ndarray",",","Dataset","]",")",":","The","train","data","needed","to","retrieve","the","training","embeddings","and","fit","the","second","VAE","in","the","latent","space",".","eval_data","(","Union","[","torch.Tensor",",","np.ndarray",",","Dataset","]",")",":","The","train","data","needed","to","retrieve","the","evaluation","embeddings","and","fit","the","second","VAE","in","the","latent","space",".","training_config","(","BaseTrainerConfig",")",":","the","training","config","to","use","to","fit","the","second","VAE",".","batch_size","(","int",")",":","The","batch","size","to","use","to","retrieve","the","embeddings",".","Default",":","64","."]
code: def fit(
      self,
      train_data: Union[torch.Tensor, np.ndarray, Dataset],
      eval_data: Union[torch.Tensor, np.ndarray, Dataset, None] = None,
      training_config: BaseTrainerConfig = None,
      batch_size: int = 64,
  ):
      """Method to fit the sampler from the training data
  
      Args:
          train_data (Union[torch.Tensor, np.ndarray, Dataset]): The train data needed to
              retrieve the training embeddings and fit the second VAE in the latent space.
          eval_data (Union[torch.Tensor, np.ndarray, Dataset]): The train data needed to retrieve
              the evaluation embeddings and fit the second VAE in the latent space.
          training_config (BaseTrainerConfig): the training config to use to fit the second VAE.
          batch_size (int): The batch size to use to retrieve the embeddings. Default: 64.
      """
  
      data_processor = DataProcessor()
      if not isinstance(train_data, Dataset):
          train_data = data_processor.process_data(train_data)
          train_dataset = data_processor.to_dataset(train_data)
  
      else:
          train_dataset = train_data
  
      train_loader = DataLoader(
          dataset=train_dataset,
          batch_size=batch_size,
          shuffle=False,
          collate_fn=collate_dataset_output,
      )
  
      z = []
  
      try:
          with torch.no_grad():
              for _, inputs in enumerate(train_loader):
                  inputs = self._set_inputs_to_device(inputs)
                  encoder_output = self.model(inputs)
                  z_ = encoder_output.z
                  z.append(z_)
  
      except RuntimeError:
          for _, inputs in enumerate(train_loader):
              inputs = self._set_inputs_to_device(inputs)
              encoder_output = self.model(inputs)
              z_ = encoder_output.z.detach()
              z.append(z_)
  
      train_data = torch.cat(z)
      train_dataset = data_processor.to_dataset(train_data)
  
      eval_dataset = None
  
      if eval_data is not None:
          if not isinstance(eval_data, Dataset):
              eval_data = data_processor.process_data(eval_data)
              eval_dataset = data_processor.to_dataset(eval_data)
  
          else:
              eval_dataset = eval_data
  
          eval_loader = DataLoader(
              dataset=eval_dataset,
              batch_size=batch_size,
              shuffle=False,
              collate_fn=collate_dataset_output,
          )
  
          z = []
  
          try:
              with torch.no_grad():
                  for _, inputs in enumerate(eval_loader):
                      inputs = self._set_inputs_to_device(inputs)
                      encoder_output = self.model(inputs)
                      z_ = encoder_output.z
                      z.append(z_)
  
          except RuntimeError:
              for _, inputs in enumerate(eval_loader):
                  inputs = self._set_inputs_to_device(inputs)
                  encoder_output = self.model(inputs)
                  z_ = encoder_output.z.detach()
                  z.append(z_)
  
          eval_data = torch.cat(z)
          eval_dataset = data_processor.to_dataset(eval_data)
  
      trainer = BaseTrainer(
          model=self.second_vae,
          train_dataset=train_dataset,
          eval_dataset=eval_dataset,
          training_config=training_config,
      )
  
      trainer.train()
  
      self.second_vae = VAE.load_from_folder(
          os.path.join(trainer.training_dir, "final_model")
      ).to(self.device)
  
      shutil.rmtree(trainer.training_dir)
  
      self.is_fitted = True
  
imports: import os
  import shutil
  from typing import Union
  import numpy
  import torch
  import torch.nn
  from torch.utils.data import DataLoader, Dataset
  from ...data.datasets import collate_dataset_output
  from ...data.preprocessors import DataProcessor
  from ...models import VAE, VAEConfig
  from ...models.base.base_utils import ModelOutput
  from ...models.nn import BaseDecoder, BaseEncoder
  from ...trainers import BaseTrainer, BaseTrainerConfig
  from ..base import BaseSampler
  from .two_stage_sampler_config import TwoStageVAESamplerConfig
label_desc: Use image node_id 2 for calling the TwoStageVAESampler obj's underlying member method code with example usage: obj.fit(train_data, eval_data, training_config, batch_size) without return types