fix: expand dim for scalar numpy when freezing tensors to IConstantLa…

…yers (#3251)
pytorch · Oct 22, 2024 · 40193dc · 40193dc
1 parent aa37194
commit 40193dc
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/py/torch_tensorrt/dynamo/conversion/converter_utils.py b/py/torch_tensorrt/dynamo/conversion/converter_utils.py
@@ -505,6 +505,13 @@ def convert_with_type_enforcement(
                         _LOGGER.debug(
                             f"Freezing tensor {name}_constant_{index} to TRT IConstantLayer"
                         )
+                        # If the candidate is a scalar Numpy array, handle it accordingly
+                        if isinstance(candidate, np.ndarray) and candidate.shape == ():
+                            _LOGGER.debug(
+                                f"Scalar numpy detected at {name}_constant_{index}, adding rank (1,)"
+                            )
+                            # Convert the scalar to a 1D array (with a single element)
+                            candidate = np.expand_dims(candidate, axis=0)
                         new_value = get_trt_tensor(
                             ctx, candidate, name + f"_constant_{index}"
                         )