fix: fix bug in torch-tensorrt 2.8.0

Author: GdoongMathew

src/lightning/pytorch/core/module.py

@@ -1661,18 +1661,26 @@ def forward(self, x):
         self.to(device)
 
         if file_path is not None:
-            if ir == "ts" and output_format != "torchscript":
-                raise ValueError(
-                    "TensorRT with IR mode 'ts' only supports output format 'torchscript'."
-                    f" The current output format is {output_format}."
+            if ir == "ts":
+                if output_format != "torchscript":
+                    raise ValueError(
+                        "TensorRT with IR mode 'ts' only supports output format 'torchscript'."
+                        f" The current output format is {output_format}."
+                    )
+                assert isinstance(trt_obj, (torch.jit.ScriptModule, torch.jit.ScriptFunction)), (
+                    f"Expected TensorRT object to be a ScriptModule, but got {type(trt_obj)}."
+                )
+                # Because of https://github.com/pytorch/TensorRT/issues/3775,
+                # we'll need to take special care for the ScriptModule
+                torch.jit.save(trt_obj, file_path)
+            else:
+                torch_tensorrt.save(
+                    trt_obj,
+                    file_path,
+                    inputs=input_sample,
+                    output_format=output_format,
+                    retrace=retrace,
                 )
-            torch_tensorrt.save(
-                trt_obj,
-                file_path,
-                inputs=input_sample,
-                output_format=output_format,
-                retrace=retrace,
-            )
         return trt_obj
 
     @_restricted_classmethod

tests/tests_pytorch/models/test_torch_tensorrt.py

@@ -130,9 +130,8 @@ def test_tensorrt_save_ir_type(ir, export_type):
     "ir",
     ["default", "dynamo", "ts"],
 )
-@RunIf(tensorrt=True, min_cuda_gpus=1, min_torch="2.2.0", max_torch="2.8.0")
+@RunIf(tensorrt=True, min_cuda_gpus=1, min_torch="2.2.0")
 def test_tensorrt_export_reload(output_format, ir, tmp_path):
-    # todo remove max_torch once https://github.com/pytorch/TensorRT/issues/3775 is fixed
     if ir == "ts" and output_format == "exported_program":
         pytest.skip("TorchScript cannot be exported as exported_program")