code: def test_save_load_checkpoint(self):
      class MyModel(objax.Module):
          def __init__(self, depth=16, width=2):
              self.model = wide_resnet.WideResNet(
                  3, 10, depth=depth, width=width
              )
              self.opt = objax.optimizer.Momentum(self.model.vars())
              self.ema = objax.optimizer.ExponentialMovingAverage(
                  self.model.vars(), momentum=0.999
              )
  
      m1 = MyModel(depth=28, width=2)
      m2 = MyModel(depth=28, width=2)
  
      # make sure that variables are different right after initialization
      var_abs_diff = 0
      v2 = m2.vars()
      for k, v in m1.vars().items():
          var_abs_diff += jn.sum(jn.abs(v.value - v2[k].value))
      self.assertGreater(var_abs_diff, 1e-3)
  
      with tempfile.TemporaryDirectory() as logdir:
          ckpt = objax.io.checkpoint.Checkpoint(logdir, keep_ckpts=2)
          ckpt.save(m1.vars(), 1)
          ckpt.save(m1.vars(), 2)
          ckpt.save(m1.vars(), 3)
          restore_idx, _ = ckpt.restore(m2.vars())
          self.assertEqual(restore_idx, 3)
  
      v2 = m2.vars()
      for k, v in m1.vars().items():
          self.assertEqual(
              v.value.tolist(),
              v2[k].value.tolist(),
              msg=f"Variable {k} value is differing.",
          )
  
      m3 = MyModel(depth=28, width=1)
  
      with tempfile.TemporaryDirectory() as logdir:
          ckpt = objax.io.checkpoint.Checkpoint(logdir, keep_ckpts=2)
          ckpt.save(m1.vars(), 1)
          with self.assertRaises(AssertionError):
              ckpt.restore(m3.vars())
  
      m3 = MyModel(depth=16, width=2)
  
      with tempfile.TemporaryDirectory() as logdir:
          ckpt = objax.io.checkpoint.Checkpoint(logdir, keep_ckpts=2)
          ckpt.save(m1.vars(), 1)
          with self.assertRaises(AssertionError):
              ckpt.restore(m3.vars())
  
file_path: objax/tests/testio.py
imports: import io
  import tempfile
  import unittest
  import jax.numpy
  import objax
  from objax.zoo import wide_resnet
label_desc: Use image node_id 1 for calling the TestCheckpoint obj's underlying member method code with example usage: obj.test_save_load_checkpoint() without return types