def test_train_save_restore_sharded(self):
    """Saves and restores a sharded checkpoint to check for equivalence."""
    if fastmath.local_device_count() < 2:
      return  # multi-accelerator only
    base.N_WEIGHTS_SHARDS = fastmath.local_device_count()
    train_data = data.Serial(lambda _: _very_simple_data(2, 2),
                             data.CountAndSkip('simple_data'))
    task = training.TrainTask(
        train_data(), tl.L2Loss(), optimizers.Adam(.0001))
    eval_task = training.EvalTask(
        _very_simple_data(2, 2),  # deliberately re-using training data
        [tl.L2Loss()],
        metric_names=['SGD.L2Loss'])
    tmp_dir = self.create_tempdir().full_path

    def _make_model_and_session():
      m = tl.Serial(tl.Dense(2))
      ts = training.Loop(m, [task], eval_tasks=[eval_task],
                         eval_at=lambda step_n: step_n % 2 == 0,
                         output_dir=tmp_dir)
      return m, ts

    _, training_session = _make_model_and_session()
    self.assertEqual(0, training_session.step)
    training_session.run(n_steps=1)
    training_session.save_checkpoint('model')
    _, training_session2 = _make_model_and_session()
    training_session2.run(n_steps=1)
    base.N_WEIGHTS_SHARDS = 1
Saves and restores a sharded checkpoint to check for equivalence.
trax/supervised/training_test.py