File size: 585 Bytes
aa8c79c
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
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
@@ -62,6 +62,8 @@ 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