Method create_test_trainer

official/core/base_trainer_test.py:175–185

(self, config, model_dir=None, task=None)

    super().tearDown()

  def create_test_trainer(self, config, model_dir=None, task=None):
    task = task or mock_task.MockTask(config.task, logging_dir=model_dir)
    ckpt_exporter = train_lib.maybe_create_best_ckpt_exporter(config, model_dir)
    trainer = trainer_lib.Trainer(
        config,
        task,
        model=task.build_model(),
        optimizer=task.create_optimizer(config.trainer.optimizer_config,
                                        config.runtime),
        checkpoint_exporter=ckpt_exporter)
    return trainer

  @combinations.generate(all_strategy_combinations())
  def test_trainer_train(self, distribution):