chore: address review comments

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

py/torch_tensorrt/dynamo/compile.py

@@ -8,6 +8,7 @@
 from torch_tensorrt import EngineCapability, Device
 from torch.fx.passes.pass_manager import PassManager
 from torch.fx.passes.shape_prop import ShapeProp
+from torch.fx.passes.splitter_base import SplitResult
 from torch_tensorrt.dynamo.aten_tracer import trace
 from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter, TRTSplitterSetting
 from torch_tensorrt.dynamo.lowering import (
@@ -17,6 +18,7 @@
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo.utils import prepare_inputs, prepare_device
 from torch_tensorrt.dynamo.backend import torch_tensorrt_backend
+from torch_tensorrt.dynamo.backend.backends import _compile_module
 from torch_tensorrt.dynamo.conversion import convert_module
 
 from torch_tensorrt.dynamo._defaults import (
@@ -124,7 +126,7 @@ def compile(
 
 
 def _compile_graph(
-    split_result: TRTSplitter,
+    split_result: SplitResult,
     inputs: Any,
     settings: CompilationSettings = CompilationSettings(),
     **kwargs,

tests/py/ts/api/test_module_fallback.py

@@ -23,6 +23,7 @@ def test_fallback_resnet18(self):
             },
             "enabled_precisions": {torch.float},
             "torch_executed_modules": ["torchvision.models.resnet.BasicBlock"],
+            "ir": "ts",
         }
         trt_mod = torchtrt.compile(self.model, **compile_spec)
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))
@@ -49,6 +50,7 @@ def test_fallback_mobilenet_v2(self):
                 "torchvision.models.mobilenetv2.ConvBNActivation"
             ],
             "min_block_size": 5,
+            "ir": "ts",
         }
         trt_mod = torchtrt.compile(self.model, **compile_spec)
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))

tests/py/ts/api/test_operator_fallback.py

@@ -23,6 +23,7 @@ def test_fallback_resnet18(self):
             },
             "enabled_precisions": {torch.float},
             "torch_executed_ops": ["aten::add"],
+            "ir": "ts",
         }
         trt_mod = torchtrt.compile(self.model, **compile_spec)
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))
@@ -49,6 +50,7 @@ def test_fallback_resnet18_with_tensor_domain(self):
             },
             "enabled_precisions": {torch.float},
             "torch_executed_ops": ["aten::add"],
+            "ir": "ts",
         }
         trt_mod = torchtrt.compile(self.model, **compile_spec)
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))
@@ -72,6 +74,7 @@ def test_fallback_mobilenet_v2(self):
             },
             "enabled_precisions": {torch.float},
             "torch_executed_ops": ["aten::hardtanh"],
+            "ir": "ts",
         }
         trt_mod = torchtrt.compile(self.model, **compile_spec)
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))