internal change

PiperOrigin-RevId: 336042908

official/core/train_lib.py

@@ -169,11 +169,9 @@ def run_experiment(distribution_strategy: tf.distribute.Strategy,
   """
 
   with distribution_strategy.scope():
-    model = task.build_model()
     trainer = train_utils.create_trainer(
         params,
         task,
-        model=model,
         model_dir=model_dir,
         train='train' in mode,
         evaluate=('eval' in mode) or run_post_eval,

official/core/train_utils.py

@@ -34,7 +34,6 @@ from official.modeling.hyperparams import config_definitions
 
 def create_trainer(params: config_definitions.ExperimentConfig,
                    task: base_task.Task,
-                   model: tf.keras.Model,
                    model_dir: str,
                    train: bool,
                    evaluate: bool,
@@ -42,6 +41,7 @@ def create_trainer(params: config_definitions.ExperimentConfig,
   """Create trainer."""
   del model_dir
   logging.info('Running default trainer.')
+  model = task.build_model()
   trainer = base_trainer.Trainer(
       params,
       task,