chore: address small bugs

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Apr 17, 2024 · a83ba8c · a83ba8c
1 parent 46ecb24
commit a83ba8c
Show file tree

Hide file tree

Showing 6 changed files with 19 additions and 7 deletions.
diff --git a/py/torch_tensorrt/__init__.py b/py/torch_tensorrt/__init__.py
@@ -103,7 +103,12 @@ def _register_with_torch() -> None:
 _register_with_torch()
 
 from torch_tensorrt._Device import Device  # noqa: F401
-from torch_tensorrt._enums import DeviceType, dtype, memory_format  # noqa: F401
+from torch_tensorrt._enums import (  # noqa: F401
+    DeviceType,
+    EngineCapability,
+    dtype,
+    memory_format,
+)
 from torch_tensorrt._Input import Input  # noqa: F401
 from torch_tensorrt.runtime import *  # noqa: F403
 

diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py
@@ -125,7 +125,7 @@ def validate_conversion(self) -> Set[str]:
 
     @staticmethod
     def _args_str(args: List[Any]) -> str:
-        def clean_repr(x: Any) -> Any:
+        def clean_repr(x: Any, depth: int = 0) -> Any:
             if isinstance(x, trt.ITensor):
                 return f"{x.name} <tensorrt.ITensor [shape={x.shape}, dtype={x.dtype}]>"
             elif isinstance(x, torch.Tensor):
@@ -134,8 +134,11 @@ def clean_repr(x: Any) -> Any:
                 return (
                     f"<torch.Tensor as np.ndarray [shape={x.shape}, dtype={x.dtype}]>"
                 )
-            elif isinstance(x, Sequence):
-                return type(x)([clean_repr(i) for i in x])  # type: ignore[call-arg]
+            elif isinstance(x, Sequence) and not isinstance(x, str):
+                if depth < 3:
+                    return type(x)([clean_repr(i, depth=depth + 1) for i in x])  # type: ignore[call-arg]
+                else:
+                    return "(...)"
             else:
                 return x
 

diff --git a/tests/py/dynamo/backend/test_backend_compiler.py b/tests/py/dynamo/backend/test_backend_compiler.py
@@ -1,11 +1,11 @@
+# type: ignore
 from copy import deepcopy
 
 import torch
+import torch_tensorrt
 from torch.testing._internal.common_utils import TestCase, run_tests
 from torch_tensorrt.dynamo.partitioning import fast_partition
 
-import torch_tensorrt
-
 from ..testing_utilities import DECIMALS_OF_AGREEMENT, lower_graph_testing
 
 

diff --git a/tests/py/dynamo/conversion/harness.py b/tests/py/dynamo/conversion/harness.py
@@ -239,6 +239,7 @@ def run_test(
         compilation_settings = CompilationSettings(
             enabled_precisions={dtype._from(precision)},
             truncate_long_and_double=True,
+            debug=True,
         )
 
         input_specs = [Input.from_tensor(i) for i in inputs]

diff --git a/tests/py/dynamo/conversion/test_pad_aten.py b/tests/py/dynamo/conversion/test_pad_aten.py
@@ -1,3 +1,4 @@
+# type: ignore
 import torch
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests

diff --git a/tests/py/ts/integrations/test_to_backend_api.py b/tests/py/ts/integrations/test_to_backend_api.py
@@ -29,7 +29,9 @@ def setUp(self):
                         "dla_core": 0,
                         "allow_gpu_fallback": True,
                     },
-                    "capability": torchtrt.ts.EngineCapability.DEFAULT,
+                    "capability": torchtrt.EngineCapability.STANDARD.to(
+                        torchtrt._C.EngineCapability
+                    ),
                     "num_avg_timing_iters": 1,
                     "disable_tf32": False,
                 }