code: def test_checkpoint_saving(self, ae, trainer, training_configs):
      dir_path = training_configs.output_dir
  
      # Make a training step
      step_1_loss = trainer.train_step(epoch=1)
  
      model = deepcopy(trainer.model)
      optimizer = deepcopy(trainer.optimizer)
  
      trainer.save_checkpoint(dir_path=dir_path, epoch=0, model=model)
  
      checkpoint_dir = os.path.join(dir_path, "checkpoint_epoch_0")
  
      assert os.path.isdir(checkpoint_dir)
  
      files_list = os.listdir(checkpoint_dir)
  
      assert set(
          ["model.pt", "optimizer.pt", "training_config.json"]
      ).issubset(set(files_list))
  
      # check pickled custom decoder
      if not ae.model_config.uses_default_decoder:
          assert "decoder.pkl" in files_list
  
      else:
          assert not "decoder.pkl" in files_list
  
      # check pickled custom encoder
      if not ae.model_config.uses_default_encoder:
          assert "encoder.pkl" in files_list
  
      else:
          assert not "encoder.pkl" in files_list
  
      model_rec_state_dict = torch.load(
          os.path.join(checkpoint_dir, "model.pt")
      )["model_state_dict"]
  
      assert all(
          [
              torch.equal(
                  model_rec_state_dict[key].cpu(),
                  model.state_dict()[key].cpu(),
              )
              for key in model.state_dict().keys()
          ]
      )
  
      # check reload full model
      model_rec = AutoModel.load_from_folder(
          os.path.join(checkpoint_dir)
      )
  
      assert all(
          [
              torch.equal(
                  model_rec.state_dict()[key].cpu(),
                  model.state_dict()[key].cpu(),
              )
              for key in model.state_dict().keys()
          ]
      )
  
      assert type(model_rec.encoder.cpu()) == type(model.encoder.cpu())
      assert type(model_rec.decoder.cpu()) == type(model.decoder.cpu())
  
      optim_rec_state_dict = torch.load(
          os.path.join(checkpoint_dir, "optimizer.pt")
      )
  
      assert all(
          [
              dict_rec == dict_optimizer
              for (dict_rec, dict_optimizer) in zip(
                  optim_rec_state_dict["param_groups"],
                  optimizer.state_dict()["param_groups"],
              )
          ]
      )
  
      assert all(
          [
              dict_rec == dict_optimizer
              for (dict_rec, dict_optimizer) in zip(
                  optim_rec_state_dict["state"],
                  optimizer.state_dict()["state"],
              )
          ]
      )
  
imports: import os
  from copy import deepcopy
  import pytest
  import torch
  from pythae.customexception import BadInheritanceError
  from pythae.models import AE, AEConfig, AutoModel
  from pythae.models.base.base_utils import ModelOutput
  from pythae.pipelines import GenerationPipeline, TrainingPipeline
  from pythae.samplers import GaussianMixtureSamplerConfig, IAFSamplerConfig, MAFSamplerConfig, NormalSamplerConfig
  from pythae.trainers import BaseTrainer, BaseTrainerConfig
  from tests.data.custom_architectures import Decoder_AE_Conv, Encoder_AE_Conv, NetBadInheritance
label_desc: Use image node_id 9 for calling the Test_AE_Training obj's underlying member method code with example usage: obj.test_checkpoint_saving(ae, trainer, training_configs) without return types