(#3) Defense: Logging images generated by defense models

betarixm · May 2, 2022 · afb63f0 · afb63f0
1 parent 1b8944d
commit afb63f0
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/defense/models.py b/src/defense/models.py
@@ -159,8 +159,10 @@ def custom_callbacks(self) -> List[keras.callbacks.Callback]:
     progress = keras.utils.Progbar(test_set.take(1).cardinality().numpy())
 
     with defense_model.tensorboard_file_writer().as_default():
-        for idx, (x, _) in enumerate(test_set.take(1)):
+        for idx, (x, _) in enumerate(test_set):
             y = defense_model.predict(x)
-            tf.summary.image(f"{args.defense.upper()} result - {idx}", [x, y])
+            tf.summary.image(
+                f"{args.defense.upper()} processing result", [x, y], step=idx
+            )
 
             progress.add(1)