code: def main():
      data = np.load(
          os.path.join(
              PATH,
              f"data/dsprites",
              "dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz",
          ),
          encoding="latin1",
      )
      train_data = torch.from_numpy(data["imgs"]).float()
  
      data_input_dim = tuple(train_data.shape[1:])
  
      model_config = BetaTCVAEConfig(
          input_dim=data_input_dim,
          latent_dim=10,
          reconstruction_loss="bce",
          beta=6,
          gamma=1,
          alpha=1,
          use_mss=False,
      )
  
      model = BetaTCVAE(
          model_config=model_config,
          encoder=Encoder(model_config),
          decoder=Decoder(model_config),
      )
  
      ### Set the training config
      training_config = BaseTrainerConfig(
          output_dir="reproducibility/dsprites",
          per_device_train_batch_size=1000,
          per_device_eval_batch_size=1000,
          num_epochs=50,
          learning_rate=1e-3,
          steps_saving=None,
          steps_predict=None,
          no_cuda=False,
          optimizer_cls="Adam",
          optimizer_params={"eps": 1e-4},
      )
  
      ### 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 = BaseTrainer(
          model=model,
          train_dataset=train_dataset,
          training_config=training_config,
          callbacks=None,
      )
  
      trainer.train()
  
file_path: benchmark_VAE/examples/scripts/reproducibility/betatcvae.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 BetaTCVAE, BetaTCVAEConfig
  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