add validator to keep static inputs

pytorch · Feb 17, 2024 · 65bbf45 · 65bbf45
1 parent 8c41f14
commit 65bbf45
Showing 1 changed file with 21 additions and 4 deletions.
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -2184,10 +2184,27 @@ def aten_ops_avg_pool(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.adaptive_avg_pool2d.default)
-@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)
+def adaptive_pool_static_input_validator(pool_node: Node) -> bool:
+    output_size = args_bounds_check(pool_node.args, 1)
+    return all([x > 0 for x in output_size])
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten.adaptive_avg_pool2d.default,
+    capability_validator=adaptive_pool_static_input_validator,
+)
+@dynamo_tensorrt_converter(
+    torch.ops.aten._adaptive_avg_pool2d.default,
+    capability_validator=adaptive_pool_static_input_validator,
+)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.adaptive_avg_pool3d.default,
+    capability_validator=adaptive_pool_static_input_validator,
+)
+@dynamo_tensorrt_converter(
+    torch.ops.aten._adaptive_avg_pool3d.default,
+    capability_validator=adaptive_pool_static_input_validator,
+)
 def aten_ops_adaptive_avg_poolNd(
     ctx: ConversionContext,
     target: Target,