diff --git a/src/mapmos/paper_pipeline.py b/src/mapmos/paper_pipeline.py
index 29f1000..8536ef3 100644
--- a/src/mapmos/paper_pipeline.py
+++ b/src/mapmos/paper_pipeline.py
@@ -158,19 +158,14 @@ def _run_pipeline(self):
                 torch.tensor(gt_labels, dtype=torch.int32),
             )
 
-            belief_labels_scan = belief_labels_with_map
-            if self.visualize:
-                belief_map = self.belief.get_belief(map_points)
-                belief_labels_map = self.model.to_label(belief_map)
-                self.visualizer.update(
-                    scan_points,
-                    map_points,
-                    pred_labels_scan,
-                    pred_labels_map,
-                    belief_labels_scan,
-                    belief_labels_map,
-                    self.odometry.last_pose,
-                )
+            self.visualizer.update(
+                scan_points,
+                map_points,
+                pred_labels_scan,
+                pred_labels_map,
+                self.belief,
+                self.odometry.last_pose,
+            )
 
             # Probabilistic volumetric fusion with scan and moving map predictions and delay
             self.buffer.append([scan_index, scan_points, gt_labels])
diff --git a/src/mapmos/pipeline.py b/src/mapmos/pipeline.py
index 89ec6df..23e2a23 100644
--- a/src/mapmos/pipeline.py
+++ b/src/mapmos/pipeline.py
@@ -174,6 +174,7 @@ def _run_pipeline(self):
 
             start_time = time.perf_counter_ns()
             self.belief.update_belief(points_stacked, logits_stacked)
+            self.belief.get_belief(scan_points)
             self.times_belief.append(time.perf_counter_ns() - start_time)
 
             self.visualizer.update(