Efficient Tensor conversion from list of numpy arrays (#1071)

jmschrei · Nov 13, 2023 · a446005 · a446005
1 parent f8ed453
commit a446005
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pomegranate/_utils.py b/pomegranate/_utils.py
@@ -56,7 +56,10 @@ def _cast_as_tensor(value, dtype=None):
 			return value
 		else:
 			return value.type(dtype)
-
+
+	if isinstance(value, list) and all(isinstance(v, numpy.ndarray) for v in value):
+		value = numpy.array(value)
+
 	if isinstance(value, (float, int, list, tuple, numpy.ndarray)):
 		if dtype is None:
 			return torch.tensor(value)