From 17a5f216498bf6038463b0ac67a9fe055b20f8ed Mon Sep 17 00:00:00 2001 From: cehongwang Date: Sun, 21 Jun 2026 22:36:33 +0000 Subject: [PATCH] fix: restore static allocation on ResourceAllocationStrategy exit The __exit__ method re-applied the same dynamically_allocate_resources value used by __enter__, so leaving the context never restored the original (static) allocation mode as documented. Set it to False on exit and remove the leftover debug print in __enter__. --- py/torch_tensorrt/dynamo/runtime/_ResourceAllocator.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/py/torch_tensorrt/dynamo/runtime/_ResourceAllocator.py b/py/torch_tensorrt/dynamo/runtime/_ResourceAllocator.py index 3e570f4d78..211bbbb068 100644 --- a/py/torch_tensorrt/dynamo/runtime/_ResourceAllocator.py +++ b/py/torch_tensorrt/dynamo/runtime/_ResourceAllocator.py @@ -21,7 +21,6 @@ def __init__( self.dynamically_allocate_resources = dynamically_allocate_resources def __enter__(self) -> None: - print("Entering resource allocator context") for name, submodule in self.compiled_module.named_modules(): if "_run_on_acc" in name: submodule.use_dynamically_allocated_resources( @@ -32,5 +31,5 @@ def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None: for name, submodule in self.compiled_module.named_modules(): if "_run_on_acc" in name: submodule.use_dynamically_allocated_resources( - dynamically_allocate_resources=self.dynamically_allocate_resources + dynamically_allocate_resources=False )