Fix for tf 2

Arno1235 · Mar 22, 2023 · db8b791 · db8b791
1 parent 19d4db0
commit db8b791
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -238,7 +238,8 @@ def train(self, data_provider, epochs, output_path, restore=False, train_summary
     def pred_image(self, image_in, checkpoint_path, num):
         #self.checkpoint = tfe.Checkpoint(model=self.model.net)
         self.checkpoint = tf.train.Checkpoint(model=self.model.net)
-        self.checkpoint.restore(checkpoint_path)
+        # self.checkpoint.restore(checkpoint_path) # TEST
+        self.checkpoint.read(checkpoint_path)
 
         preds = []
         for icounter in range(num):