chore: move functions to organize code better

pytorch · May 22, 2024 · af31bce · af31bce
1 parent c4965cd
commit af31bce
Showing 1 changed file with 34 additions and 34 deletions.
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -800,6 +800,40 @@ def aten_ops_tile(
     )
 
 
+def zero_output_validator(node: Node) -> bool:
+    if 0 in node.args[1]:
+        _LOGGER.debug(
+            f"We do not support output tensor {node.args[1]} tensors with zero-sized dimensions for this operation."
+        )
+        return False
+    else:
+        return True
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten.as_strided.default,
+    capability_validator=zero_output_validator,
+)
+@dynamo_tensorrt_converter(torch.ops.aten.as_strided.default)
+def aten_ops_as_strided(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.slice.as_strided(
+        ctx,
+        target,
+        source_ir=SourceIR.ATEN,
+        name=name,
+        input=args[0],
+        size=args[1],
+        stride=args[2],
+        storage_offset=args_bounds_check(args, 3, None),
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.permute.default)
 @enforce_tensor_types(
     {
@@ -2186,40 +2220,6 @@ def aten_ops_linear(
     )
 
 
-def zero_output_validator(node: Node) -> bool:
-    if 0 in node.args[1]:
-        _LOGGER.debug(
-            f"We do not support output tensor {node.args[1]} tensors with zero-sized dimensions for this operation."
-        )
-        return False
-    else:
-        return True
-
-
-@dynamo_tensorrt_converter(
-    torch.ops.aten.as_strided.default,
-    capability_validator=zero_output_validator,
-)
-@dynamo_tensorrt_converter(torch.ops.aten.as_strided.default)
-def aten_ops_as_strided(
-    ctx: ConversionContext,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.slice.as_strided(
-        ctx,
-        target,
-        source_ir=SourceIR.ATEN,
-        name=name,
-        input=args[0],
-        size=args[1],
-        stride=args[2],
-        storage_offset=args_bounds_check(args, 3, None),
-    )
-
-
 def avg_pool_param_validator(pool_node: Node) -> bool:
     ceil_mode = args_bounds_check(pool_node.args, 4, False)
     divisor_override = args_bounds_check(pool_node.args, 6)