update tests and enforce input type

pytorch · Apr 10, 2024 · f1b1ad0 · f1b1ad0
1 parent 61b74f3
commit f1b1ad0
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -2185,6 +2185,11 @@ def aten_ops_avg_pool(
 
 
 @dynamo_tensorrt_converter(torch.ops.aten.adaptive_avg_pool1d.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
 def aten_ops_adaptive_avg_pool1d(
     ctx: ConversionContext,
     target: Target,
@@ -2206,6 +2211,11 @@ def aten_ops_adaptive_avg_pool1d(
 @dynamo_tensorrt_converter(torch.ops.aten._adaptive_avg_pool2d.default)
 @dynamo_tensorrt_converter(torch.ops.aten.adaptive_avg_pool3d.default)
 @dynamo_tensorrt_converter(torch.ops.aten._adaptive_avg_pool3d.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
 def aten_ops_adaptive_avg_poolNd(
     ctx: ConversionContext,
     target: Target,

diff --git a/tests/py/dynamo/conversion/test_adaptive_avgpool_aten.py b/tests/py/dynamo/conversion/test_adaptive_avgpool_aten.py
@@ -76,7 +76,6 @@ def forward(self, x):
         self.run_test(
             TestModule(),
             inputs,
-            # use_dynamo_tracer=True,
             enable_passes=True,
         )
 
@@ -155,7 +154,6 @@ def forward(self, x):
         self.run_test(
             TestModule(),
             inputs,
-            # use_dynamo_tracer=True,
             enable_passes=True,
         )
 
@@ -268,7 +266,6 @@ def forward(self, x):
         self.run_test(
             TestModule(),
             inputs,
-            # use_dynamo_tracer=True,
             enable_passes=True,
         )