code: def main():
      train_data = (
          np.load(os.path.join(PATH, f"data/celeba", "train_data.npz"))[
              "data"
          ]
          / 255.0
      )
      eval_data = (
          np.load(os.path.join(PATH, f"data/celeba", "eval_data.npz"))[
              "data"
          ]
          / 255.0
      )
  
      data_input_dim = tuple(train_data.shape[1:])
  
      model_config = RAE_GP_Config(
          input_dim=data_input_dim,
          latent_dim=16,
          embedding_weight=1e-3,
          reg_weight=1e-6,
      )
  
      model = RAE_GP(
          model_config=model_config,
          encoder=Encoder(model_config),
          decoder=Decoder(model_config),
      )
  
      ### Set training config
      training_config = BaseTrainerConfig(
          output_dir="my_models_on_celeba",
          per_device_train_batch_size=100,
          per_device_eval_batch_size=100,
          num_epochs=100,
          learning_rate=0.001,
          steps_saving=None,
          steps_predict=100,
          no_cuda=False,
          scheduler_cls="ReduceLROnPlateau",
          scheduler_params={
              "factor": 0.5,
              "patience": 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()
  
file_path: benchmark_VAE/examples/scripts/reproducibility/rae_gp.py
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 RAE_GP, RAE_GP_Config
  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 1 for calling a global function with example usage: main() without return types