Fix chesapeake plot test

microsoft · Aug 9, 2024 · 919bbb2 · 919bbb2
1 parent dd27fd4
commit 919bbb2
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tests/datasets/test_chesapeake.py b/tests/datasets/test_chesapeake.py
@@ -225,6 +225,9 @@ def test_plot(self, dataset: ChesapeakeCVPR) -> None:
         plt.close()
         dataset.plot(x, show_titles=False)
         plt.close()
-        x['prediction'] = x['mask'][:, :, 0].clone().unsqueeze(2)
+        if x['mask'].ndim == 2:
+            x['prediction'] = x['mask'].clone()
+        else:
+            x['prediction'] = x['mask'][0, :, :].clone()
         dataset.plot(x)
         plt.close()