Addressed the comments

Author: cehongwang

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -188,11 +188,12 @@ def _pack_engine_info(self) -> List[str | bytes]:
         engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = str(
             int(self.requires_output_allocator)
         )
-        print(f"PROVIDED RESOURCE ALLOCATION STRATEGY: {self.dynamically_allocate_resources}")
+        logger.info(
+            f"PROVIDED RESOURCE ALLOCATION STRATEGY: {self.dynamically_allocate_resources}"
+        )
         engine_info[RESOURCE_ALLOCATION_STRATEGY_IDX] = str(
             int(self.dynamically_allocate_resources)
         )
-        print(engine_info[RESOURCE_ALLOCATION_STRATEGY_IDX])
 
         return engine_info
 
@@ -221,9 +222,13 @@ def set_device_memory_budget(self, budget_bytes: int) -> int:
     def _reset_captured_graph(self) -> None:
         self.engine.reset_captured_graph()
 
-    def use_dynamically_allocated_resources(self, dynamically_allocate_resources: bool = False) -> None:
+    def use_dynamically_allocated_resources(
+        self, dynamically_allocate_resources: bool = False
+    ) -> None:
         self.dynamically_allocate_resources = dynamically_allocate_resources
-        self.engine.use_dynamically_allocated_resources(self.dynamically_allocate_resources)
+        self.engine.use_dynamically_allocated_resources(
+            self.dynamically_allocate_resources
+        )
 
     def setup_engine(self) -> None:
         """

tests/py/dynamo/runtime/test_005_dynamic_allocation.py

@@ -2,18 +2,29 @@
 import unittest
 
 import torch
+import torch.nn.functional as F
 import torch_tensorrt as torch_trt
+from torch import nn
 from torch.testing._internal.common_utils import TestCase
 from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity
 
 assertions = unittest.TestCase()
 
-if importlib.util.find_spec("torchvision"):
-    import torchvision.models as models
-
 
 class TestDynamicAllocation(TestCase):
     def test_dynamic_allocation(self):
+
+        class net(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.conv1 = nn.Conv2d(3, 6, 3, 1)
+                self.conv2 = nn.Conv2d(6, 16, 3, 1)
+
+            def forward(self, x):
+                x = F.relu(self.conv1(x))
+                x = F.relu(self.conv2(x))
+                return x
+
         inputs = [torch.rand((100, 3, 224, 224)).to("cuda")]
 
         settings = {
@@ -25,7 +36,7 @@ def test_dynamic_allocation(self):
             "dynamically_allocate_resources": True,
         }
 
-        model = models.resnet152(pretrained=True).eval().to("cuda")
+        model = net().eval().to("cuda")
         compiled_module = torch_trt.compile(model, inputs=inputs, **settings)
         compiled_module(*inputs)