chore: small linting

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

Hide file tree

Showing 4 changed files with 5 additions and 11 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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()