diff --git a/src/data_gradients/dataset_adapters/formatters/segmentation.py b/src/data_gradients/dataset_adapters/formatters/segmentation.py
index 3087217d..7f2438d2 100644
--- a/src/data_gradients/dataset_adapters/formatters/segmentation.py
+++ b/src/data_gradients/dataset_adapters/formatters/segmentation.py
@@ -81,7 +81,7 @@ def _format_images(self, images: Tensor) -> Tensor:
         return images
 
     def _format_labels(self, labels: Tensor) -> Tensor:
-        labels = labels.squeeze((1, -1))  # If (BS, 1, H, W) or (BS, H, W, 1) -> (BS, H, W)
+        labels = labels.squeeze(1).squeeze(-1)  # If (BS, 1, H, W) or (BS, H, W, 1) -> (BS, H, W)
         if labels.ndim == 3:
             labels = ensure_hard_labels(labels, n_classes=self.data_config.get_n_classes(), threshold_value=self.threshold_value)
         elif labels.ndim == 4:
diff --git a/src/data_gradients/feature_extractors/object_detection/bounding_boxes_iou.py b/src/data_gradients/feature_extractors/object_detection/bounding_boxes_iou.py
index 970a8ecc..f03fa1b9 100644
--- a/src/data_gradients/feature_extractors/object_detection/bounding_boxes_iou.py
+++ b/src/data_gradients/feature_extractors/object_detection/bounding_boxes_iou.py
@@ -65,6 +65,7 @@ def _compute_cumulative_counts_at_thresholds(self, df: pd.DataFrame, class_names
         return counts
 
     def aggregate(self) -> Feature:
+
         df = pd.DataFrame(self.data).sort_values(by="class_id")
 
         bins = np.linspace(0, 1, self.num_bins + 1)
@@ -96,8 +97,7 @@ def aggregate(self) -> Feature:
         xticklabels = [f"IoU < {bins[x]:.2f}" for x in range(1, len(bins))]
 
         if not data:
-            self._show_plot = False
-            return Feature(data=None, plot_options=None, json={})
+            return Feature(data=None, plot_options=None, json={}, title="Intersection of Bounding Boxes", description="No Overlap between any class")
 
         # Height of the plot is proportional to the number of classes
         figsize_x = min(max(10, len(bins)), 25)
diff --git a/src/data_gradients/feature_extractors/segmentation/classes_heatmap_per_class.py b/src/data_gradients/feature_extractors/segmentation/classes_heatmap_per_class.py
index 19594b8d..9a82ea5e 100644
--- a/src/data_gradients/feature_extractors/segmentation/classes_heatmap_per_class.py
+++ b/src/data_gradients/feature_extractors/segmentation/classes_heatmap_per_class.py
@@ -31,7 +31,7 @@ def update(self, sample: SegmentationSample):
             self.class_names = sample.class_names
 
         # (H, W) -> (C, H, W)
-        n_classes = np.max(sample.mask)
+        n_classes = np.max(list(sample.class_names.keys()))
         mask_onehot = mask_to_onehot(mask_categorical=sample.mask, n_classes=n_classes)
         mask_onehot = mask_onehot.transpose((1, 2, 0))  # H, W, C -> C, H, W
 
diff --git a/src/data_gradients/managers/abstract_manager.py b/src/data_gradients/managers/abstract_manager.py
index 557df4a9..8b4ad86b 100644
--- a/src/data_gradients/managers/abstract_manager.py
+++ b/src/data_gradients/managers/abstract_manager.py
@@ -172,15 +172,16 @@ def post_process(self, interrupted=False):
                 else:
                     warning = feature.warning
 
-                section.add_feature(
-                    FeatureSummary(
-                        name=feature.title,
-                        description=self._format_feature_description(feature.description),
-                        image_path=image_path,
-                        warning=warning,
-                        notice=feature.notice,
+                if not feature_error:
+                    section.add_feature(
+                        FeatureSummary(
+                            name=feature.title,
+                            description=self._format_feature_description(feature.description),
+                            image_path=image_path,
+                            warning=warning,
+                            notice=feature.notice,
+                        )
                     )
-                )
             summary.add_section(section)
 
         print("Dataset successfully analyzed!")