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: Main sampling function of the sampler.
  
  Args:
      num_samples (int): The number of samples to generate
      batch_size (int): The batch size to use during sampling
      output_dir (str): The directory where the images will be saved. If does not exist the
          folder is created. If None: the images are not saved. Defaults: None.
      return_gen (bool): Whether the sampler should directly return a tensor of generated
          data. Default: True.
      save_sampler_config (bool): Whether to save the sampler config. It is saved in
          output_dir
  
  Returns:
      ~torch.Tensor: The generated images
docstr_tok: ["Main","sampling","function","of","the","sampler",".","Args",":","num_samples","(","int",")",":","The","number","of","samples","to","generate","batch_size","(","int",")",":","The","batch","size","to","use","during","sampling","output_dir","(","str",")",":","The","directory","where","the","images","will","be","saved",".","If","does","not","exist","the","folder","is","created",".","If","None",":","the","images","are","not","saved",".","Defaults",":","None",".","return_gen","(","bool",")",":","Whether","the","sampler","should","directly","return","a","tensor","of","generated","data",".","Default",":","True",".","save_sampler_config","(","bool",")",":","Whether","to","save","the","sampler","config",".","It","is","saved","in","output_dir","Returns",":","~torch.Tensor",":","The","generated","images"]
code: def sample(
      self,
      num_samples: int = 1,
      batch_size: int = 500,
      output_dir: str = None,
      return_gen: bool = True,
      save_sampler_config: bool = False,
  ) -> torch.Tensor:
      """Main sampling function of the sampler.
  
      Args:
          num_samples (int): The number of samples to generate
          batch_size (int): The batch size to use during sampling
          output_dir (str): The directory where the images will be saved. If does not exist the
              folder is created. If None: the images are not saved. Defaults: None.
          return_gen (bool): Whether the sampler should directly return a tensor of generated
              data. Default: True.
          save_sampler_config (bool): Whether to save the sampler config. It is saved in
              output_dir
  
      Returns:
          ~torch.Tensor: The generated images
      """
  
      if not self.is_fitted:
          raise ArithmeticError(
              "The sampler needs to be fitted by calling sampler.fit() method"
              "before sampling."
          )
  
      full_batch_nbr = int(num_samples / batch_size)
      last_batch_samples_nbr = num_samples % batch_size
  
      x_gen_list = []
  
      for i in range(full_batch_nbr):
          u = torch.randn(batch_size, self.model.latent_dim).to(
              self.device
          )
          z = self.second_vae.decoder(u).reconstruction
          x_gen = self.model.decoder(z)["reconstruction"].detach()
  
          if output_dir is not None:
              for j in range(batch_size):
                  self.save_img(
                      x_gen[j],
                      output_dir,
                      "%08d.png" % int(batch_size * i + j),
                  )
  
          x_gen_list.append(x_gen)
  
      if last_batch_samples_nbr > 0:
          u = torch.randn(
              last_batch_samples_nbr, self.model.latent_dim
          ).to(self.device)
          z = self.second_vae.decoder(u).reconstruction
          x_gen = self.model.decoder(z)["reconstruction"].detach()
  
          if output_dir is not None:
              for j in range(last_batch_samples_nbr):
                  self.save_img(
                      x_gen[j],
                      output_dir,
                      "%08d.png" % int(batch_size * full_batch_nbr + j),
                  )
  
          x_gen_list.append(x_gen)
  
      if save_sampler_config:
          self.save(output_dir)
  
      if return_gen:
          return torch.cat(x_gen_list, dim=0)
  
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 3 for calling the TwoStageVAESampler obj's underlying member method code with example usage: obj.sample(num_samples, batch_size, output_dir, return_gen, save_sampler_config) and returns: torch