Move to_numpy implementation to converter_util

Author: apbose

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -543,3 +543,27 @@ def type_cast(
     layer_i.set_output_type(0, cast_type)
     set_layer_name(layer_i, target, f"{name}_dtype_change")
     return layer_i.get_output(0)
+
+
+def to_numpy(tensor: Optional[torch.Tensor]) -> Optional[np.ndarray]:
+    """
+    Convert a PyTorch Tensor to a Numpy Array. If the tensor is
+    quantized it will be dequantized first.
+
+    Args:
+        tensor (Optional[torch.Tensor]): A PyTorch tensor or None.
+
+    Returns:
+        A Numpy array.
+    """
+
+    if tensor is None:
+        return tensor
+
+    assert isinstance(
+        tensor, torch.Tensor
+    ), f"to_numpy can only be called on None or a torch.Tensor, got: {tensor}"
+    if tensor.is_quantized:
+        tensor = tensor.dequantize()
+
+    return tensor.cpu().detach().contiguous().numpy()

py/torch_tensorrt/fx/converters/operator.py

@@ -22,6 +22,7 @@
 from .converter_utils import prepend_ones
 from .converter_utils import has_dynamic_shape
 from .converter_utils import get_shape_with_dynamic_shape
+from .converter_utils import to_numpy
 
 from ..types import (
     Shape,
@@ -278,30 +279,6 @@ def trunc_div(
     return output
 
 
-def to_numpy(tensor: Optional[torch.Tensor]) -> Optional[np.ndarray]:
-    """
-    Convert a PyTorch Tensor to a Numpy Array. If the tensor is
-    quantized it will be dequantized first.
-
-    Args:
-        tensor (Optional[torch.Tensor]): A PyTorch tensor or None.
-
-    Returns:
-        A Numpy array.
-    """
-
-    if tensor is None:
-        return tensor
-
-    assert isinstance(
-        tensor, torch.Tensor
-    ), f"to_numpy can only be called on None or a torch.Tensor, got: {tensor}"
-    if tensor.is_quantized:
-        tensor = tensor.dequantize()
-
-    return tensor.cpu().detach().contiguous().numpy()
-
-
 def trt_dtype_to_torch_dtype(trt_dtype):
     table = {
         trt.bool: torch.bool,