From 2d1d036c07a6b069fa2acf5a51b6503c491e839c Mon Sep 17 00:00:00 2001
From: Dobromir Popov <dobromir.popov@gateway.one>
Date: Mon, 17 Nov 2025 21:12:20 +0200
Subject: [PATCH] More T train fixes

---
 ANNOTATE/core/real_training_adapter.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/ANNOTATE/core/real_training_adapter.py b/ANNOTATE/core/real_training_adapter.py
index b95748b..a16dca3 100644
--- a/ANNOTATE/core/real_training_adapter.py
+++ b/ANNOTATE/core/real_training_adapter.py
@@ -2128,13 +2128,16 @@ class RealTrainingAdapter:
             # MEMORY FIX: Final cleanup
             logger.info("   Final memory cleanup...")
             
-            # Clear cached batches
-            for batch in cached_batches:
-                for key in list(batch.keys()):
-                    if isinstance(batch[key], torch.Tensor):
-                        del batch[key]
-            cached_batches.clear()
-            del cached_batches
+            # Clear grouped batches (cached_batches was already cleared earlier)
+            # Note: Don't delete batch contents as they may be referenced elsewhere
+            # Just clear the list reference - Python GC will handle cleanup
+            try:
+                if grouped_batches:
+                    grouped_batches.clear()
+                    del grouped_batches
+            except NameError:
+                # grouped_batches already cleaned up or doesn't exist
+                pass
             
             gc.collect()
             if torch.cuda.is_available():