chore: fix some tests

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

Author: peri044

py/torch_tensorrt/dynamo/_refit.py

@@ -274,9 +274,15 @@ def refit_module_weights(
     else:
         for name, submodule in compiled_module.named_children():
             if not isinstance(
-                submodule, (PythonTorchTensorRTModule, TorchTensorRTModule)
+                submodule,
+                (
+                    PythonTorchTensorRTModule,
+                    TorchTensorRTModule,
+                    torch.nn.modules.module.Module,
+                ),
             ):
                 continue
+
             settings = submodule.settings
 
     assert settings is not None

tests/py/dynamo/models/test_model_refit.py

@@ -540,8 +540,8 @@ def test_refit_one_engine_inline_runtime_with_weightmap():
     min_block_size = 1
     use_python_runtime = False
 
-    exp_program = torch.export.export(model, tuple(inputs))
-    exp_program2 = torch.export.export(model2, tuple(inputs))
+    exp_program = torch.export.export(model, tuple(inputs), strict=False)
+    exp_program2 = torch.export.export(model2, tuple(inputs), strict=False)
 
     trt_gm = torchtrt.dynamo.compile(
         exp_program,
@@ -551,8 +551,10 @@ def test_refit_one_engine_inline_runtime_with_weightmap():
         min_block_size=min_block_size,
         immutable_weights=False,
     )
-    torchtrt.save(trt_gm, trt_ep_path)
+    torchtrt.save(trt_gm, trt_ep_path, arg_inputs=inputs)
+
     trt_gm = torch.export.load(trt_ep_path)
+
     new_trt_gm = refit_module_weights(
         compiled_module=trt_gm,
         new_weight_module=exp_program2,
@@ -906,7 +908,7 @@ def test_refit_one_engine_inline_runtime_without_weightmap():
         min_block_size=min_block_size,
         immutable_weights=False,
     )
-    torchtrt.save(trt_gm, trt_ep_path)
+    torchtrt.save(trt_gm, trt_ep_path, arg_inputs=inputs)
     trt_gm = torch.export.load(trt_ep_path)
     new_trt_gm = refit_module_weights(
         compiled_module=trt_gm,

tests/py/dynamo/runtime/test_002_lazy_engine_init.py

@@ -314,7 +314,9 @@ def test_lazy_engine_init_cpp_serialization(self):
         trt_mod = torchtrt.compile(model, **compile_spec)
 
         with tempfile.TemporaryDirectory() as tmpdir:
-            torch_tensorrt.save(trt_mod, os.path.join(tmpdir, "tmp_trt_mod.ep"))
+            torch_tensorrt.save(
+                trt_mod, os.path.join(tmpdir, "tmp_trt_mod.ep"), arg_inputs=(input,)
+            )
             new_trt_mod = torch.export.load(os.path.join(tmpdir, "tmp_trt_mod.ep"))
 
         loaded_trt_mod = new_trt_mod.module()