fix: expand dim for scalar numpy when freezing tensors to IConstantLayers (#3251)

chohk88 · web-flow · commit 40193dcfbaac · 2024-10-22T22:51:51.000+09:00
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}"
                         )

Original file line number	Diff line number	Diff line change
`@@ -505,6 +505,13 @@ def convert_with_type_enforcement(`
`505`	`505`	`_LOGGER.debug(`
`506`	`506`	`f"Freezing tensor {name}_constant_{index} to TRT IConstantLayer"`
`507`	`507`	`)`
	`508`	`+ # If the candidate is a scalar Numpy array, handle it accordingly`
	`509`	`+ if isinstance(candidate, np.ndarray) and candidate.shape == ():`
	`510`	`+ _LOGGER.debug(`
	`511`	`+ f"Scalar numpy detected at {name}_constant_{index}, adding rank (1,)"`
	`512`	`+ )`
	`513`	`+ # Convert the scalar to a 1D array (with a single element)`
	`514`	`+ candidate = np.expand_dims(candidate, axis=0)`
`508`	`515`	`new_value = get_trt_tensor(`
`509`	`516`	`ctx, candidate, name + f"_constant_{index}"`
`510`	`517`	`)`