chore: update hw_compat

Author: peri044

tests/py/dynamo/runtime/gen_hw_compat.py

@@ -0,0 +1,33 @@
+# This script is used to generate hw_compat.ts file that's used in test_hw_compat.py
+# Generate the model on a different hardware compared to the one you're testing on to
+# verify HW compatibility feature.
+
+import torch
+import torch_tensorrt
+
+
+class MyModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(3, 16, 3, stride=1, bias=True)
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x):
+        out = self.conv(x)
+        out = self.relu(out)
+        return out
+
+
+model = MyModule().eval().cuda()
+inputs = torch.randn((1, 3, 224, 224)).to("cuda")
+
+trt_gm = torch_tensorrt.compile(
+    model,
+    ir="dynamo",
+    inputs=inputs,
+    min_block_size=1,
+    hardware_compatible=True,
+    version_compatible=True,
+)
+trt_script_model = torch.jit.trace(trt_gm, inputs)
+torch.jit.save(trt_script_model, "hw_compat.ts")

tests/py/dynamo/runtime/test_hw_compat.py

@@ -59,16 +59,14 @@ def forward(self, x):
         "Detected incorrect ABI version, please update this test case",
     )
     def test_hw_compat_3080_build(self):
-        inputs = [torch.randn(5, 7).cuda()]
+        inputs = [torch.randn(1, 3, 224, 224).cuda()]
 
         cwd = os.getcwd()
         os.chdir(os.path.dirname(os.path.realpath(__file__)))
         model = torch.jit.load("../../ts/models/hw_compat.ts").cuda()
         out = model(*inputs)
         self.assertTrue(
-            isinstance(out, tuple)
-            and len(out) == 1
-            and isinstance(out[0], torch.Tensor),
+            len(out) == 1 and isinstance(out, torch.Tensor),
             "Invalid output detected",
         )
         os.chdir(cwd)