(FIX) Fixes Lightning-Universe#430, predict batches should now be sho…

…wn for image classification.
fstroth · Jun 21, 2021 · cd3faf8 · cd3faf8
1 parent de3f851
commit cd3faf8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flash/image/classification/data.py b/flash/image/classification/data.py
@@ -136,9 +136,9 @@ def _show_images_and_labels(self, data: List[Any], num_samples: int, title: str)
         for i, ax in enumerate(axs.ravel()):
             # unpack images and labels
             if isinstance(data, list):
-                _img, _label = data[i][DefaultDataKeys.INPUT], data[i][DefaultDataKeys.TARGET]
+                _img, _label = data[i][DefaultDataKeys.INPUT], data[i].get(DefaultDataKeys.TARGET, "")
             elif isinstance(data, dict):
-                _img, _label = data[DefaultDataKeys.INPUT][i], data[DefaultDataKeys.TARGET][i]
+                _img, _label = data[DefaultDataKeys.INPUT][i], data.get(DefaultDataKeys.TARGET, [""]*(i+1))[i]
             else:
                 raise TypeError(f"Unknown data type. Got: {type(data)}.")
             # convert images to numpy