Internal change

PiperOrigin-RevId: 397851196

Internal change
PiperOrigin-RevId: 397851196
23339995 · A. Unique TensorFlower · bed0e3de · 23339995
Commit 23339995 authored Sep 20, 2021 by A. Unique TensorFlower
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

official/vision/beta/tasks/image_classification.py official/vision/beta/tasks/image_classification.py +1 -1

No files found.
--- a/official/vision/beta/tasks/image_classification.py
+++ b/official/vision/beta/tasks/image_classification.py
@@ -62,7 +62,7 @@ class ImageClassificationTask(base_task.Task):

    # Restoring checkpoint.
    if self.task_config.init_checkpoint_modules == 'all':
-      ckpt = tf.train.Checkpoint(**model.checkpoint_items)
+      ckpt = tf.train.Checkpoint(model=model)
      status = ckpt.read(ckpt_dir_or_file)
      status.expect_partial().assert_existing_objects_matched()
    elif self.task_config.init_checkpoint_modules == 'backbone':