chore: fix tests

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

Author: peri044

tests/py/dynamo/models/test_export_kwargs_serde.py

@@ -76,7 +76,7 @@ def forward(self, x, b=5, c=None, d=None):
 
     # Save the module
     trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
-    torchtrt.save(trt_gm, trt_ep_path)
+    torchtrt.save(trt_gm, trt_ep_path, retrace=False)
     # Clean up model env
     torch._dynamo.reset()
 
@@ -138,7 +138,7 @@ def forward(self, x, b=5, c=None, d=None):
 
     # Save the module
     trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
-    torchtrt.save(trt_gm, trt_ep_path)
+    torchtrt.save(trt_gm, trt_ep_path, retrace=False)
     # Clean up model env
     torch._dynamo.reset()
 
@@ -209,7 +209,7 @@ def forward(self, x, b=5, c=None, d=None):
 
     # Save the module
     trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
-    torchtrt.save(trt_gm, trt_ep_path)
+    torchtrt.save(trt_gm, trt_ep_path, retrace=False)
     # Clean up model env
     torch._dynamo.reset()
 
@@ -299,7 +299,7 @@ def forward(self, x, b=None, c=None, d=None, e=[]):
     )
     # Save the module
     trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
-    torchtrt.save(trt_gm, trt_ep_path)
+    torchtrt.save(trt_gm, trt_ep_path, retrace=False)
     # Clean up model env
     torch._dynamo.reset()
 
@@ -389,7 +389,7 @@ def forward(self, x, b=None, c=None, d=None, e=[]):
     )
     # Save the module
     trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
-    torchtrt.save(trt_gm, trt_ep_path)
+    torchtrt.save(trt_gm, trt_ep_path, retrace=False)
     # Clean up model env
     torch._dynamo.reset()
 

tests/py/dynamo/models/test_export_serde.py

@@ -56,7 +56,7 @@ def forward(self, x):
 
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     # Check Pyt and TRT exported program outputs
@@ -111,7 +111,7 @@ def forward(self, x):
 
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     # Check Pyt and TRT exported program outputs
@@ -170,7 +170,7 @@ def forward(self, x):
 
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     # Check Pyt and TRT exported program outputs
@@ -232,7 +232,7 @@ def forward(self, x):
 
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     outputs_pyt = model(input)
@@ -279,7 +279,7 @@ def test_resnet18(ir):
 
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     outputs_pyt = model(input)
@@ -331,7 +331,7 @@ def test_resnet18_cpu_offload(ir):
             msg="Model should be offloaded to CPU",
         )
         model.cuda()
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     outputs_pyt = model(input)
@@ -380,7 +380,7 @@ def test_resnet18_dynamic(ir):
 
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
     # TODO: Enable this serialization issues are fixed
     # deser_trt_module = torchtrt.load(trt_ep_path).module()
     outputs_pyt = model(input)
@@ -413,7 +413,7 @@ def test_resnet18_torch_exec_ops_serde(ir):
 
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     outputs_pyt = deser_trt_module(input)
     outputs_trt = trt_module(input)
@@ -463,7 +463,7 @@ def forward(self, x):
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
 
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     outputs_pyt = model(input)
@@ -525,7 +525,7 @@ def forward(self, x):
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
     model.cuda()
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     outputs_pyt = model(input)
@@ -584,7 +584,7 @@ def forward(self, x):
     exp_program = torchtrt.dynamo.trace(model, **compile_spec)
     trt_module = torchtrt.dynamo.compile(exp_program, **compile_spec)
 
-    torchtrt.save(trt_module, trt_ep_path, arg_inputs=[input], retrace=False)
+    torchtrt.save(trt_module, trt_ep_path, retrace=False)
 
     deser_trt_module = torchtrt.load(trt_ep_path).module()
     outputs_pyt = model(input)