fixed data loading

gisilvs · Jul 17, 2020 · 13d4a9b · 13d4a9b
1 parent b3d80f3
commit 13d4a9b
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/train.py b/train.py
@@ -204,12 +204,12 @@ def test_step(images, targets):
     test_metric.reset_states()
 
     # training loop
-    for images, targets, true in train_ds:
-        train_step(images, targets, true)
+    for images, targets in train_ds:
+        train_step(images, targets)
 
     # test loop
-    for test_images, test_labels, test_true in test_ds:
-        test_step(test_images, test_labels, test_true)
+    for test_images, test_targets in test_ds:
+        test_step(test_images, test_targets)
 
     # log losses and metrics
     with summary_writer.as_default():