From 203afb5f2aa8e56f4a794f83a4d051e421d2516f Mon Sep 17 00:00:00 2001
From: reaperhammer <87021961+reaperhammer@users.noreply.github.com>
Date: Sat, 9 May 2026 15:40:57 +1200
Subject: [PATCH] Improve CUDA graph memory cleanup in Engine class

- __del__: Add hasattr checks before del, iterate tensors dict to delete
  each tensor individually, also clean up inputs/outputs
- reset: Same improvements plus reset cuda_graph_instance and graph
  to None to prevent stale state on reuse

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 trt_utilities.py | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/trt_utilities.py b/trt_utilities.py
index a41d338..4a12561 100644
--- a/trt_utilities.py
+++ b/trt_utilities.py
@@ -151,10 +151,20 @@ class Engine:
             except:
                 pass
 
-        del self.engine
-        del self.context
-        del self.buffers
-        del self.tensors
+        if hasattr(self, 'engine'):
+            del self.engine
+        if hasattr(self, 'context'):
+            del self.context
+        if hasattr(self, 'tensors'):
+            for key in list(self.tensors.keys()):
+                del self.tensors[key]
+            del self.tensors
+        if hasattr(self, 'buffers'):
+            del self.buffers
+        if hasattr(self, 'inputs'):
+            del self.inputs
+        if hasattr(self, 'outputs'):
+            del self.outputs
 
     def reset(self, engine_path=None):
         # Clean up CUDA graph resources first
@@ -175,10 +185,12 @@ class Engine:
             del self.engine
         if hasattr(self, 'context') and self.context is not None:
             del self.context
+        if hasattr(self, 'tensors'):
+            for key in list(self.tensors.keys()):
+                del self.tensors[key]
+            del self.tensors
         if hasattr(self, 'buffers'):
             del self.buffers
-        if hasattr(self, 'tensors'):
-            del self.tensors
 
         self.engine = None
         self.context = None
@@ -188,6 +200,8 @@ class Engine:
         self.tensors = OrderedDict()
         self.inputs = {}
         self.outputs = {}
+        self.cuda_graph_instance = None
+        self.graph = None
 
     def build(
         self,