class_docstr: Fits a Gaussian Mixture in the Autoencoder's latent space.
  
  Args:
      model (BaseAE): The vae model to sample from.
      sampler_config (BaseSamplerConfig): An instance of BaseSamplerConfig in which any sampler's
          parameters is made available. If None a default configuration is used. Default: None.
  
  .. note::
  
      The method :class:`~pythae.samplers.GaussianMixtureSampler.fit` must be called to fit the sampler
      before sampling.
class_docstr_tok: ["Fits","a","Gaussian","Mixture","in","the","Autoencoder","'s","latent","space",".","Args",":","model","(","BaseAE",")",":","The","vae","model","to","sample","from",".","sampler_config","(","BaseSamplerConfig",")",":","An","instance","of","BaseSamplerConfig","in","which","any","sampler's","parameters","is","made","available",".","If","None","a","default","configuration","is","used",".","Default",":","None",".","..","note",":",":","The","method",":","class",":","`","~pythae.samplers.GaussianMixtureSampler.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 mixture in the latent space.
      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","mixture","in","the","latent","space",".","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],
      batch_size: int = 64,
      **kwargs,
  ):
      """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 mixture in the latent space.
          batch_size (int): The batch size to use to retrieve the embeddings. Default: 64.
      """
      self.is_fitted = True
  
      if not isinstance(train_data, Dataset):
          data_processor = DataProcessor()
          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)
                  z_ = self.model(inputs).z
                  z.append(z_)
  
      except RuntimeError:
          for _, inputs in enumerate(train_loader):
              inputs = self._set_inputs_to_device(inputs)
              z_ = self.model(inputs).z.detach()
              z.append(z_)
  
      z = torch.cat(z)
  
      if self.n_components > z.shape[0]:
          self.n_components = z.shape[0]
          logger.warning(
              f"Setting the number of component to {z.shape[0]} since"
              "n_components > n_samples when fitting the gmm"
          )
  
      gmm = mixture.GaussianMixture(
          n_components=self.n_components,
          covariance_type="full",
          max_iter=2000,
          verbose=0,
          tol=1e-3,
      )
      gmm.fit(z.cpu().detach())
  
      self.gmm = gmm
  
file_path: benchmark_VAE/src/pythae/samplers/gaussian_mixture/gaussian_mixture_sampler.py
imports: import logging
  from typing import Union
  import numpy
  import torch
  from sklearn import mixture
  from torch.utils.data import DataLoader, Dataset
  from ...data.datasets import collate_dataset_output
  from ...data.preprocessors import DataProcessor
  from ...models import BaseAE
  from ..base import BaseSampler
  from .gaussian_mixture_config import GaussianMixtureSamplerConfig
label_desc: Use image node_id 2 for calling the GaussianMixtureSampler obj's underlying member method code with example usage: obj.fit(train_data, batch_size) without return types