code: def main():
      train_data = np.loadtxt(
          os.path.join(
              PATH, f"data/binary_mnist", "binarized_mnist_train.amat"
          )
      )
      eval_data = np.loadtxt(
          os.path.join(
              PATH, f"data/binary_mnist", "binarized_mnist_valid.amat"
          )
      )
      test_data = np.loadtxt(
          os.path.join(
              PATH, f"data/binary_mnist", "binarized_mnist_test.amat"
          )
      )
  
      data_input_dim = tuple(train_data.shape[1:])
  
      ### Build model
      model_config = VAMPConfig(
          input_dim=data_input_dim,
          latent_dim=40,
          reconstruction_loss="bce",
          number_components=500,
          linear_scheduling_steps=100,
      )
  
      model = VAMP(
          model_config=model_config,
          encoder=Encoder(model_config),
          decoder=Decoder(model_config),
      )
  
      ### Set training config
      training_config = BaseTrainerConfig(
          output_dir="reproducibility/binary_mnist",
          per_device_train_batch_size=100,
          per_device_eval_batch_size=100,
          num_epochs=2000,
          learning_rate=1e-4,
          steps_saving=None,
          steps_predict=None,
          no_cuda=False,
          optimizer_cls="RMSprop",
          scheduler_cls="MultiStepLR",
          scheduler_params={
              "milestones": [200, 350, 500, 750, 1000],
              "gamma": 10 ** (-1 / 5),
              "verbose": True,
          },
      )
  
      ### Process data
      data_processor = DataProcessor()
      logger.info("Preprocessing train data...")
      train_data = data_processor.process_data(train_data)
      train_dataset = data_processor.to_dataset(train_data)
  
      logger.info("Preprocessing eval data...\n")
      eval_data = data_processor.process_data(eval_data)
      eval_dataset = data_processor.to_dataset(eval_data)
  
      seed = 123
      torch.manual_seed(seed)
      torch.cuda.manual_seed_all(seed)
  
      logger.info("Using Base Trainer\n")
      trainer = BaseTrainer(
          model=model,
          train_dataset=train_dataset,
          eval_dataset=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)
  
      ### 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 pythae.data.preprocessors import DataProcessor
  from pythae.models import VAMP, AutoModel, VAMPConfig
  from pythae.models.base.base_utils import ModelOutput
  from pythae.models.nn import BaseDecoder, BaseEncoder
  from pythae.trainers import BaseTrainer, BaseTrainerConfig
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types