Fixed the comments

Author: cehongwang

examples/dynamo/mutable_torchtrt_module_example.py

@@ -65,7 +65,7 @@
 # Saving Mutable Torch TensorRT Module
 # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-# Currently, saving is only when "use_python" = False in settings
+# Currently, saving is only enabled when "use_python_runtime" = False in settings
 torch_trt.MutableTorchTensorRTModule.save(mutable_module, "mutable_module.pkl")
 reload = torch_trt.MutableTorchTensorRTModule.load("mutable_module.pkl")
 
@@ -201,24 +201,18 @@ def forward(self, a, b, c={}):
 from torch_tensorrt.dynamo._defaults import TIMING_CACHE_PATH
 
 model = models.resnet18(pretrained=True).eval().to("cuda")
-enabled_precisions = {torch.float}
-debug = False
-min_block_size = 1
-use_python_runtime = True
 
 times = []
 start = torch.cuda.Event(enable_timing=True)
 end = torch.cuda.Event(enable_timing=True)
 
-
 example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),)
-# Mark the dim0 of inputs as dynamic
 model = torch_trt.MutableTorchTensorRTModule(
     model,
-    use_python_runtime=use_python_runtime,
-    enabled_precisions=enabled_precisions,
-    debug=debug,
-    min_block_size=min_block_size,
+    use_python_runtime=True,
+    enabled_precisions={torch.float},
+    debug=True,
+    min_block_size=1,
     immutable_weights=False,
     cache_built_engines=True,
     reuse_cached_engines=True,

py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py

@@ -228,20 +228,21 @@ def forward(a, b, c=0, d=0):
         seq_len = torch.export.Dim("seq_len", min=1, max=10)
         args_dynamic_shape = ({0: seq_len}, {}) # b does not have a dynamic shape
         kwargs_dynamic_shape = {'c': {0, seq_len}, 'd': {}} # d does not have a dynamic shape
+        set_expected_dynamic_shape_range(args_dynamic_shape, kwargs_dynamic_shape)
         # Later when you call the function
         forward(*(a, b), **{c:..., d:...})
 
-
+        Reference: https://pytorch.org/docs/stable/export.html#expressing-dynamism
         Arguments:
             args_dynamic_shape (tuple[dict[Any, Any]]): Dynamic shape hint for the arg_inputs,
             kwargs_dynamic_shape: (dict[str, Any]): Dynamic shape hint for the kwarg_inputs
         """
         assert isinstance(
             args_dynamic_shape, tuple
-        ), "args dynamic shape has to be a tuple"
+        ), f"args dynamic shape has to be a tuple, but got {type(args_dynamic_shape)}"
         assert isinstance(
             kwargs_dynamic_shape, dict
-        ), "args dynamic shape has to be a dictionary"
+        ), f"args dynamic shape has to be a dictionary, but got {type(kwargs_dynamic_shape)}"
         self.kwarg_dynamic_shapes = kwargs_dynamic_shape
         self.arg_dynamic_shapes = args_dynamic_shape
 

tests/py/dynamo/runtime/test_mutable_torchtrt_module.py

@@ -66,11 +66,11 @@ def test_check_input_shape_dynamic():
     dynamic_shape = {"a": {1: dim}, "b": [{}, {}], "c": {"a": {}, "b": [{}, {}]}}
     assertions.assertFalse(
         torch_trt.MutableTorchTensorRTModule._check_inputs_shape(a, b),
-        msg=f"test_check_output_equal is not correct.",
+        msg=f"test_check_input_shape_dynamic is not correct.",
     )
     assertions.assertTrue(
         torch_trt.MutableTorchTensorRTModule._check_inputs_shape(a, b, dynamic_shape),
-        msg=f"test_check_output_equal is not correct.",
+        msg=f"test_check_input_shape_dynamic is not correct.",
     )