Added the string case

Author: cehongwang

py/torch_tensorrt/dynamo/_settings.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass, field
 from typing import Any, Collection, Optional, Set, Tuple, Union
 
-import torch
 from torch.fx.node import Target
 from torch_tensorrt._Device import Device
 from torch_tensorrt._enums import EngineCapability, dtype
@@ -145,14 +144,19 @@ class CompilationSettings:
     offload_module_to_cpu: bool = OFFLOAD_MODULE_TO_CPU
 
     def __getstate__(self) -> dict[str, Any]:
+        from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
+            ConverterRegistry,
+        )
+
         state = self.__dict__.copy()
-        state["torch_executed_ops"] = {str(op) for op in state["torch_executed_ops"]}
+        state["torch_executed_ops"] = {
+            op if isinstance(op, str) else ConverterRegistry.qualified_name_or_str(op)
+            for op in state["torch_executed_ops"]
+        }
         return state
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         self.__dict__.update(state)
-        ops_str = self.torch_executed_ops
-        self.torch_executed_ops = {getattr(torch.ops, op) for op in ops_str}
 
 
 _SETTINGS_TO_BE_ENGINE_INVARIANT = (

tests/py/dynamo/models/test_models.py

@@ -84,6 +84,41 @@ def test_resnet18_cpu_offload(ir):
     torch._dynamo.reset()
 
 
+@pytest.mark.unit
+def test_resnet18_torch_exec_ops(ir):
+    model = models.resnet18(pretrained=True).eval().to("cuda")
+    input = torch.randn((1, 3, 224, 224)).to("cuda")
+
+    compile_spec = {
+        "inputs": [
+            torchtrt.Input(
+                min_shape=(1, 3, 224, 224),
+                opt_shape=(8, 3, 224, 224),
+                max_shape=(16, 3, 224, 224),
+                dtype=torch.float32,
+            )
+        ],
+        "ir": ir,
+        "enabled_precisions": {torch.float32, torch.float16, torch.bfloat16},
+        "min_block_size": 1,
+        "debug": True,
+        "output_format": "exported_program",
+        "cache_built_engines": True,
+        "reuse_cached_engines": True,
+        "torch_executed_ops": {torch.ops.aten.matmul, "torch.ops.aten.add"},
+    }
+
+    trt_mod = torchtrt.compile(model, **compile_spec)
+    cos_sim = cosine_similarity(model(input), trt_mod(input))
+    assertions.assertTrue(
+        cos_sim > COSINE_THRESHOLD,
+        msg=f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+    )
+
+    # Clean up model env
+    torch._dynamo.reset()
+
+
 @pytest.mark.unit
 def test_mobilenet_v2(ir):
     model = models.mobilenet_v2(pretrained=True).eval().to("cuda")