code: def main():
      train_data = (
          np.load(os.path.join(PATH, f"data/celeba", "train_data.npz"))[
              "data"
          ]
          / 255.0
      )
  
      data_input_dim = tuple(train_data.shape[1:])
  
      model_config = Adversarial_AE_Config(
          input_dim=data_input_dim,
          latent_dim=64,
          reconstruction_loss="mse",
          adversarial_loss_scale=0.5,
          reconstruction_loss_scale=0.05,
          deterministic_posterior=True,
      )
  
      model = Adversarial_AE(
          model_config=model_config,
          encoder=Encoder(model_config),
          decoder=Decoder(model_config),
          discriminator=Discriminator(model_config),
      )
  
      ### Set training config
      training_config = AdversarialTrainerConfig(
          output_dir="my_models_on_celeba",
          per_device_train_batch_size=100,
          per_device_eval_batch_size=100,
          num_epochs=100,
          autoencoder_learning_rate=3e-4,
          discriminator_learning_rate=1e-3,
          steps_saving=3,
          steps_predict=1000,
          no_cuda=False,
          autoencoder_scheduler_cls="LambdaLR",
          autoencoder_scheduler_params={
              "lr_lambda": lambda epoch: 1 * (epoch < 30)
              + 0.5 * (30 <= epoch < 50)
              + 0.2 * (50 <= epoch),
              "verbose": True,
          },
          discriminator_scheduler_cls="LambdaLR",
          discriminator_scheduler_params={
              "lr_lambda": lambda epoch: 1 * (epoch < 30)
              + 0.5 * (30 <= epoch < 50)
              + 0.2 * (50 <= epoch),
              "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)
  
      seed = 123
      torch.manual_seed(seed)
      torch.cuda.manual_seed_all(seed)
  
      logger.info("Using Base Trainer\n")
      trainer = AdversarialTrainer(
          model=model,
          train_dataset=train_dataset,
          training_config=training_config,
          callbacks=None,
      )
  
      trainer.train()
  
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 Adversarial_AE, Adversarial_AE_Config
  from pythae.models.base.base_utils import ModelOutput
  from pythae.models.nn import BaseDecoder, BaseDiscriminator, BaseEncoder
  from pythae.trainers import AdversarialTrainer, AdversarialTrainerConfig
label_desc: Use image node_id 1 for calling a global function with example usage: main() without return types