diff --git a/utils/plots.py b/utils/plots.py
index db6f94a..5901ca2 100644
--- a/utils/plots.py
+++ b/utils/plots.py
@@ -59,7 +59,9 @@ def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detec
     n:              Maximum number of feature maps to plot
     save_dir:       Directory to save results
     """
-    if 'Detect' not in module_type:
+    if ('Detect'
+            not in module_type) and ('Segment'
+                                     not in module_type):  # 'Detect' for Object Detect task,'Segment' for Segment task
         batch, channels, height, width = x.shape  # batch, channels, height, width
         if height > 1 and width > 1:
             f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png"  # filename