diff --git a/cutler/demo/predictor.py b/cutler/demo/predictor.py | |
index 0920159..6fc8053 100644 | |
--- a/cutler/demo/predictor.py | |
+++ b/cutler/demo/predictor.py | |
class VisualizationDemo(object): | |
) | |
if "instances" in predictions: | |
instances = predictions["instances"].to(self.cpu_device) | |
+ if self.metadata.name == '__unused': | |
+ instances.remove('pred_classes') | |
vis_output = visualizer.draw_instance_predictions(predictions=instances) | |
return predictions, vis_output | |