Merge pull request #404 from aashish-kumar/patch-2

Visualization fix handling multi-channel images
cleverhans-lab · May 29, 2018 · 81fa7bb · 81fa7bb
2 parents 0e481f2 + a6268d7
commit 81fa7bb
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/cleverhans/utils.py b/cleverhans/utils.py
@@ -161,6 +161,10 @@ def pair_visual(original, adversarial, figure=None):
     """
     import matplotlib.pyplot as plt
 
+    # Squeeze the image to remove single-dimensional entries from array shape
+    original = np.squeeze(original)
+    adversarial = np.squeeze(adversarial)
+
     # Ensure our inputs are of proper shape
     assert(len(original.shape) == 2 or len(original.shape) == 3)