Fix TensorRT tensor address binding - add explicit set_tensor_address for all tensors before inference

2026-02-27 02:32:38 +01:00
parent c19731ec58
commit 0c0f4b2d10
1 changed files with 5 additions and 0 deletions
@@ -445,6 +445,11 @@ class Engine:
                self.tensors[name].copy_(buf)

            nvtx.range_push("TensorRT Inference")
+            
+            # Set tensor addresses explicitly for all bindings
+            for name, tensor in self.tensors.items():
+                self.context.set_tensor_address(name, tensor.data_ptr())
+            
            if use_cuda_graph:
                if self.cuda_graph_instance is not None:
                    CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream.ptr))