Fix conversion from MLX to torch tensor

dottxt-ai · Dec 9, 2024 · fb43f4f · fb43f4f
1 parent 147b03e
commit fb43f4f
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/outlines/processors/base_logits_processor.py b/outlines/processors/base_logits_processor.py
@@ -107,9 +107,12 @@ def _to_torch(tensor_like: Array) -> torch.Tensor:
             return torch.tensor(tensor_like)
 
         elif is_mlx_array_type(type(tensor_like)):
-            # mlx -> torch -> mlx conversion docs:
-            # https://ml-explore.github.io/mlx/build/html/usage/numpy.html
-            return torch.from_dlpack(tensor_like)
+            import mlx.core as mx
+
+            # https://ml-explore.github.io/mlx/build/html/usage/numpy.html#pytorch
+            return torch.from_dlpack(
+                np.array(tensor_like.astype(mx.float32), copy=False)
+            )
 
         elif is_jax_array_type(type(tensor_like)):
             import jax