code: def train_dataset():
      return torch.load(
          os.path.join(PATH, "data/mnist_clean_train_dataset_sample")
      )
  
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 1 for calling a global function with example usage: train_dataset() and returns: torch