class_docstr: Importance Weighted Autoencoder model.
  
  Args:
      model_config (IWAEConfig): The IWAE configuration setting the main
          parameters of the model.
  
      encoder (BaseEncoder): An instance of BaseEncoder (inheriting from `torch.nn.Module` which
          plays the role of encoder. This argument allows you to use your own neural networks
          architectures if desired. If None is provided, a simple Multi Layer Preception
          (https://en.wikipedia.org/wiki/Multilayer_perceptron) is used. Default: None.
  
      decoder (BaseDecoder): An instance of BaseDecoder (inheriting from `torch.nn.Module` which
          plays the role of decoder. This argument allows you to use your own neural networks
          architectures if desired. If None is provided, a simple Multi Layer Preception
          (https://en.wikipedia.org/wiki/Multilayer_perceptron) is used. Default: None.
  
  .. note::
      For high dimensional data we advice you to provide you own network architectures. With the
      provided MLP you may end up with a ``MemoryError``.
docstr: The VAE model
  
  Args:
      inputs (BaseDataset): The training dataset with labels
  
  Returns:
      ModelOutput: An instance of ModelOutput containing all the relevant parameters
docstr_tok: ["The","VAE","model","Args",":","inputs","(","BaseDataset",")",":","The","training","dataset","with","labels","Returns",":","ModelOutput",":","An","instance","of","ModelOutput","containing","all","the","relevant","parameters"]
code: def forward(self, inputs: BaseDataset, **kwargs):
      """
      The VAE model
  
      Args:
          inputs (BaseDataset): The training dataset with labels
  
      Returns:
          ModelOutput: An instance of ModelOutput containing all the relevant parameters
  
      """
  
      x = inputs["data"]
  
      encoder_output = self.encoder(x)
  
      mu, log_var = (
          encoder_output.embedding,
          encoder_output.log_covariance,
      )
  
      mu = mu.unsqueeze(1).repeat(1, self.n_samples, 1)
      log_var = log_var.unsqueeze(1).repeat(1, self.n_samples, 1)
  
      std = torch.exp(0.5 * log_var)
  
      z, _ = self._sample_gauss(mu, std)
  
      recon_x = self.decoder(z.reshape(-1, self.latent_dim))[
          "reconstruction"
      ].reshape(x.shape[0], self.n_samples, -1)
  
      loss, recon_loss, kld = self.loss_function(
          recon_x, x, mu, log_var, z
      )
  
      output = ModelOutput(
          recon_loss=recon_loss,
          reg_loss=kld,
          loss=loss,
          recon_x=recon_x.reshape(x.shape[0], self.n_samples, -1)[
              :, 0, :
          ].reshape_as(x),
          z=z.reshape(x.shape[0], self.n_samples, -1)[:, 0, :].reshape(
              -1, self.latent_dim
          ),
      )
  
      return output
  
file_path: benchmark_VAE/src/pythae/models/iwae/iwae_model.py
imports: from typing import Optional
  import torch
  import torch.nn.functional
  from ...data.datasets import BaseDataset
  from ..base.base_utils import ModelOutput
  from ..nn import BaseDecoder, BaseEncoder
  from ..vae import VAE
  from .iwae_config import IWAEConfig
label_desc: Use image node_id 2 for calling the IWAE obj's underlying member method code with example usage: obj.forward(inputs) and returns: output