fix bug in test.py (#819)

c2701467 · Ziyi Wu · GitHub · 31ef559a · c2701467
Unverified Commit c2701467 authored Aug 02, 2021 by Ziyi Wu Committed by GitHub Aug 02, 2021
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

tools/test.py tools/test.py +2 -1

No files found.
--- a/tools/test.py
+++ b/tools/test.py
@@ -181,7 +181,8 @@ def main():
    # palette for visualization in segmentation tasks
    if 'PALETTE' in checkpoint.get('meta', {}):
        model.PALETTE = checkpoint['meta']['PALETTE']
-    else:
+    elif hasattr(dataset, 'PALETTE'):
+        # segmentation dataset has `PALETTE` attribute
        model.PALETTE = dataset.PALETTE

    if not distributed: