From be0e0e3c7167b96e502bb7c77e4671d8499c4a90 Mon Sep 17 00:00:00 2001
From: Naren Dasan <naren@narendasan.com>
Date: Wed, 17 Apr 2024 13:54:29 -0700
Subject: [PATCH] chore: small linting

Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>
---
 docsrc/py_api/torch_tensorrt.rst      | 6 +-----
 py/torch_tensorrt/_enums.py           | 3 +--
 py/torch_tensorrt/dynamo/_compiler.py | 6 +++---
 py/torch_tensorrt/ts/_Device.py       | 1 -
 4 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/docsrc/py_api/torch_tensorrt.rst b/docsrc/py_api/torch_tensorrt.rst
index 22fda13ba2..eb8285e103 100644
--- a/docsrc/py_api/torch_tensorrt.rst
+++ b/docsrc/py_api/torch_tensorrt.rst
@@ -37,10 +37,6 @@ Classes
    :members:
    :special-members: __init__
 
-.. autoclass:: TRTModuleNext
-   :members:
-   :special-members: __init__
-
 Enums
 -------
 
@@ -50,7 +46,7 @@ Enums
 
 .. autoclass:: EngineCapability
 
-.. autoclass:: TensorFormat
+.. autoclass:: memory_format
 
 Submodules
 ----------
diff --git a/py/torch_tensorrt/_enums.py b/py/torch_tensorrt/_enums.py
index 724a2a9346..350d8a299e 100644
--- a/py/torch_tensorrt/_enums.py
+++ b/py/torch_tensorrt/_enums.py
@@ -299,7 +299,6 @@ def try_to(
         use_default: bool,
     ) -> Optional[Union[torch.dtype, trt.DataType, np.dtype, dtype]]:
         try:
-            print(self)
             casted_format = self.to(t, use_default)
             return casted_format
         except (ValueError, TypeError) as e:
@@ -689,7 +688,7 @@ def to(
             else:
                 raise ValueError("Provided an unsupported engine capability")
 
-        elif t == DeviceType:
+        elif t == EngineCapability:
             return self
 
         elif ENABLED_FEATURES.torchscript_frontend:
diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py
index bdea78c12c..ed9a0bb7ae 100644
--- a/py/torch_tensorrt/dynamo/_compiler.py
+++ b/py/torch_tensorrt/dynamo/_compiler.py
@@ -419,9 +419,9 @@ def convert_module_to_trt_engine(
     module: torch.fx.GraphModule,
     method_name: str = "forward",
     inputs: Optional[Sequence[Input | torch.Tensor]] = None,
-    enabled_precisions: Set[torch.dtype | dtype] | Tuple[torch.dtype | dtype] = (
-        dtype.float32,
-    ),
+    enabled_precisions: (
+        Set[torch.dtype | dtype] | Tuple[torch.dtype | dtype]
+    ) = _defaults.ENABLED_PRECISIONS,
     debug: bool = _defaults.DEBUG,
     workspace_size: int = _defaults.WORKSPACE_SIZE,
     min_block_size: int = _defaults.MIN_BLOCK_SIZE,
diff --git a/py/torch_tensorrt/ts/_Device.py b/py/torch_tensorrt/ts/_Device.py
index b1829d494b..3ae10a9c4d 100644
--- a/py/torch_tensorrt/ts/_Device.py
+++ b/py/torch_tensorrt/ts/_Device.py
@@ -51,7 +51,6 @@ def __init__(self, *args: Any, **kwargs: Any):
             - Device(gpu_id=1)
         """
         super().__init__(*args, **kwargs)
-        print(self)
 
     def _to_internal(self) -> _C.Device:
         internal_dev = _C.Device()