diff --git a/geo_deep_learning/tools/callbacks/segmentation_callback.py b/geo_deep_learning/tools/callbacks/segmentation_callback.py
index a9eaf4b3..83b67027 100644
--- a/geo_deep_learning/tools/callbacks/segmentation_callback.py
+++ b/geo_deep_learning/tools/callbacks/segmentation_callback.py
@@ -76,8 +76,4 @@ def _log_visualizations(self, trainer, batch_idx):
                                                         artifact_file = artifact_file,
                                                         run_id=trainer.logger.run_id)
             except Exception as e:
-                print(f"Error in logging visualizations: {e}")
-            
-            finally:
-                self.current_batch = None
-                self.current_outputs = None
\ No newline at end of file
+                print(f"Error in logging visualizations: {e}")
\ No newline at end of file