code: def main():
      ### Load data
      train_data = torch.tensor(
          np.load(os.path.join(PATH, f"data/mnist", "train_data.npz"))[
              "data"
          ]
          / 255.0
      )
      eval_data = torch.tensor(
          np.load(os.path.join(PATH, f"data/mnist", "eval_data.npz"))[
              "data"
          ]
          / 255.0
      )
  
      train_data = torch.cat((train_data, eval_data))
  
      test_data = (
          np.load(os.path.join(PATH, f"data/mnist", "test_data.npz"))[
              "data"
          ]
          / 255.0
      )
      data_input_dim = tuple(train_data.shape[1:])
  
      model_config = PIWAEConfig(
          input_dim=data_input_dim,
          latent_dim=50,
          reconstruction_loss="bce",
          number_samples=8,
          number_gradient_estimates=8,
      )
  
      model = PIWAE(
          model_config=model_config,
          encoder=Encoder(model_config),
          decoder=Decoder(model_config),
      )
  
      ### Set training config
      training_config = CoupledOptimizerTrainerConfig(
          output_dir="reproducibility/mnist",
          per_device_train_batch_size=20,
          per_device_eval_batch_size=20,
          num_epochs=3280,
          learning_rate=1e-3,
          steps_saving=1000,
          steps_predict=None,
          no_cuda=False,
          encoder_optimizer_cls="Adam",
          encoder_optimizer_params={"eps": 1e-4},
          decoder_optimizer_cls="Adam",
          decoder_optimizer_params={"eps": 1e-4},
          encoder_scheduler_cls="MultiStepLR",
          encoder_scheduler_params={
              "milestones": [2, 5, 14, 28, 41, 122, 365, 1094],
              "gamma": 10 ** (-1 / 7),
              "verbose": True,
          },
          decoder_scheduler_cls="MultiStepLR",
          decoder_scheduler_params={
              "milestones": [2, 5, 14, 28, 41, 122, 365, 1094],
              "gamma": 10 ** (-1 / 7),
              "verbose": True,
          },
      )
  
      logger.info("Preprocessing train data...")
      train_dataset = DynBinarizedMNIST(train_data)
  
      seed = 123
      torch.manual_seed(seed)
      torch.cuda.manual_seed_all(seed)
  
      trainer = CoupledOptimizerTrainer(
          model=model,
          train_dataset=train_dataset,
          eval_dataset=None,  # eval_dataset,
          training_config=training_config,
          callbacks=None,
      )
  
      trainer.train()
  
      ### Reload model
      trained_model = (
          AutoModel.load_from_folder(
              os.path.join(
                  training_config.output_dir,
                  f"{trainer.model.model_name}_training_{trainer._training_signature}",
                  "final_model",
              )
          )
          .to(device)
          .eval()
      )
  
      test_data = torch.tensor(test_data).to(device).type(torch.float)
      test_data = (
          test_data
          > torch.distributions.Uniform(0, 1)
          .sample(test_data.shape)
          .to(test_data.device)
      ).float()
  
      ### Compute NLL
      with torch.no_grad():
          nll = []
          for i in range(5):
              nll_i = trained_model.get_nll(
                  test_data, n_samples=5000, batch_size=5000
              )
              logger.info(f"Round {i+1} nll: {nll_i}")
              nll.append(nll_i)
  
      logger.info(f"\nmean_nll: {np.mean(nll)}")
      logger.info(f"\std_nll: {np.std(nll)}")
  
imports: import logging
  import os
  from typing import List
  import numpy
  import torch
  import torch.nn
  from torch.utils.data import Dataset
  from pythae.data.datasets import DatasetOutput
  from pythae.models import PIWAE, AutoModel, PIWAEConfig
  from pythae.models.base.base_utils import ModelOutput
  from pythae.models.nn import BaseDecoder, BaseEncoder
  from pythae.trainers import CoupledOptimizerTrainer, CoupledOptimizerTrainerConfig
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types