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","."]
code: def __init__(
      self, model: VAE, sampler_config: TwoStageVAESamplerConfig = None
  ):
      assert issubclass(model.__class__, VAE), (
          "The TwoStageVAESampler is only"
          f"applicable for VAE based models. Got {model.__class__}."
      )
  
      self.is_fitted = False
  
      if sampler_config is None:
          sampler_config = TwoStageVAESamplerConfig()
  
      BaseSampler.__init__(
          self, model=model, sampler_config=sampler_config
      )
  
      second_vae_config = VAEConfig(
          latent_dim=model.model_config.latent_dim,
          reconstruction_loss=sampler_config.reconstruction_loss,
      )
  
      self.second_vae = VAE(
          model_config=second_vae_config,
          encoder=SecondEncoder(model, sampler_config),
          decoder=SecondDecoder(model, sampler_config),
      )
  
      self.second_vae.model_name = "Second_Stage_VAE"
  
      self.second_vae.to(self.device)
  
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 1 to create a new TwoStageVAESampler object from inherited base classes: BaseSampler with example: obj = TwoStageVAESampler(model, sampler_config)