Repository: pytorch/TensorRT

Issue:
🐛 [Bug] Compiling Models with Int-to-Tensor Casted Inputs
##  Bug Description

When compiling models with 0D tensor inputs, the following error is encountered:

```python
ERROR: [Torch-TensorRT] - 3: [executionContext.cpp::setInputShape::2259] Error Code 3: API Usage Error (Parameter check failed at: runtime/api/executionContext.cpp::setInputShape::2259, condition: engineDims.nbDims == dims.nbDims
)
```

## To Reproduce

```python
        class Multiply(torch.nn.Module):
            def forward(self, x):
                return x * 7

        inputs = [
            torch.tensor(
                3,
            ).cuda().int(),
        ]
```

## Expected behavior

The models should successfully compile.

## Environment

> Build information about Torch-TensorRT can be found by turning on debug messages

 - Torch-TensorRT Version (e.g. 1.0.0): 8c62fca7649ac2e8e74ab7d9d66a20a164e2af51
 - PyTorch Version (e.g. 1.0): `2.1.0.dev20230803+cu121`

## Additional context

Very relevant to diffusion and detectron-style models.


Discussion:

Resolving patch:
diff --git a/py/torch_tensorrt/_Input.py b/py/torch_tensorrt/_Input.py
--- a/py/torch_tensorrt/_Input.py
+++ b/py/torch_tensorrt/_Input.py
@@ -32,11 +32,11 @@ class _ShapeMode(Enum):
     shape: Optional[
         Tuple[int, ...] | Dict[str, Tuple[int, ...]]
     ] = None  #: Either a single Tuple or a dict of tuples defining the input shape. Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form ``{ "min_shape": Tuple, "opt_shape": Tuple, "max_shape": Tuple }``
-    dtype: _enums.dtype = (  # type: ignore[name-defined]
+    dtype: _enums.dtype = (
         _enums.dtype.unknown
     )  #: The expected data type of the input tensor (default: torch_tensorrt.dtype.float32)
     _explicit_set_dtype: bool = False
-    format: _enums.TensorFormat = (  # type: ignore[name-defined]
+    format: _enums.TensorFormat = (
         _enums.TensorFormat.contiguous
     )  #: The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
 
@@ -208,7 +208,7 @@ def _supported_input_size_type(input_size: Any) -> bool:
             return False
 
     @staticmethod
-    def _parse_dtype(dtype: Any) -> _enums.dtype:  # type: ignore[name-defined]
+    def _parse_dtype(dtype: Any) -> _enums.dtype:
         if isinstance(dtype, torch.dtype):
             if dtype == torch.long:
                 return _enums.dtype.long
@@ -236,7 +236,7 @@ def _parse_dtype(dtype: Any) -> _enums.dtype:  # type: ignore[name-defined]
             )
 
     @staticmethod
-    def _to_torch_dtype(dtype: _enums.dtype) -> torch.dtype:  # type: ignore[name-defined]
+    def _to_torch_dtype(dtype: _enums.dtype) -> torch.dtype:
         if dtype == _enums.dtype.long:
             return torch.long
         elif dtype == _enums.dtype.int32:
@@ -255,7 +255,7 @@ def is_trt_dtype(self) -> bool:
         return bool(self.dtype != _enums.dtype.long)
 
     @staticmethod
-    def _parse_format(format: Any) -> _enums.TensorFormat:  # type: ignore[name-defined]
+    def _parse_format(format: Any) -> _enums.TensorFormat:
         if isinstance(format, torch.memory_format):
             if format == torch.contiguous_format:
                 return _enums.TensorFormat.contiguous
@@ -337,9 +337,9 @@ def from_tensor(
             A Input object.
         """
         if not (
-            t.is_contiguous(memory_format=torch.contiguous_format)
+            disable_memory_format_check
+            or t.is_contiguous(memory_format=torch.contiguous_format)
             or t.is_contiguous(memory_format=torch.channels_last)
-            or disable_memory_format_check
         ):
             raise ValueError(
                 "Tensor does not have a supported memory format, supported formats are contiguous or channel_last"
@@ -347,8 +347,8 @@ def from_tensor(
         frmt = (
             torch.contiguous_format
             if (
-                t.is_contiguous(memory_format=torch.contiguous_format)
-                or disable_memory_format_check
+                disable_memory_format_check
+                or t.is_contiguous(memory_format=torch.contiguous_format)
             )
             else torch.channels_last
         )


Test patch:
diff --git a/tests/py/dynamo/backend/test_specialized_models.py b/tests/py/dynamo/backend/test_specialized_models.py
--- a/tests/py/dynamo/backend/test_specialized_models.py
+++ b/tests/py/dynamo/backend/test_specialized_models.py
@@ -1,7 +1,7 @@
-from utils import lower_graph_testing
-from torch.testing._internal.common_utils import run_tests, TestCase
 import torch
 import torch_tensorrt
+from torch.testing._internal.common_utils import TestCase, run_tests
+from utils import lower_graph_testing
 
 
 class TestFakeTensors(TestCase):
@@ -118,5 +118,43 @@ def forward(self, x):
         torch._dynamo.reset()
 
 
+class Test0DTensors(TestCase):
+    def test_0D_input(self):
+        class Tensor0DInput(torch.nn.Module):
+            def forward(self, x):
+                return x * 7
+
+        inputs = [
+            torch.tensor(
+                3,
+            )
+            .cuda()
+            .int(),
+        ]
+
+        fx_graph = torch.fx.symbolic_trace(Tensor0DInput())
+
+        # Validate that the results between Torch and Torch-TRT are similar
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_compile",
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
+        )
+        optimized_model_results = optimized_model(*inputs).detach().cpu()
+        torch_model_results = fx_graph(*inputs).detach().cpu()
+
+        max_diff = float(
+            torch.max(torch.abs(optimized_model_results - torch_model_results))
+        )
+        self.assertAlmostEqual(
+            max_diff,
+            0,
+            msg=f"0D-Tensor TRT outputs don't match with the original model.",
+        )
+        torch._dynamo.reset()
+
+
 if __name__ == "__main__":
     run_tests()