code: def model_sample():
      return AE(AEConfig(input_dim=(1, 28, 28)))
  
imports: import os
  from copy import deepcopy
  import pytest
  import torch
  from pythae.customexception import ModelError
  from pythae.models import AE, RHVAE, VAE, AEConfig, RHVAEConfig, VAEConfig
  from pythae.trainers import BaseTrainer, BaseTrainerConfig
  from tests.data.custom_architectures import *
label_desc: Use image node_id 2 for calling a global function with example usage: model_sample() and returns: AE