Fix coverage

microsoft · Nov 9, 2024 · 2d281ce · 2d281ce
1 parent f30e697
commit 2d281ce
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/datasets/test_vhr10.py b/tests/datasets/test_vhr10.py
@@ -85,7 +85,7 @@ def test_plot(self, dataset: VHR10) -> None:
                 x['prediction_labels'] = x['class']
                 x['prediction_boxes'] = x['bbox_xyxy']
                 x['prediction_scores'] = torch.Tensor([scores[i]])
-                if 'masks' in x:
+                if 'mask' in x:
                     x['prediction_masks'] = x['mask']
                     dataset.plot(x, show_feats='masks')
                     plt.close()