zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-22 11:14:32 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix eval in trace bugs (#612)

Browse Source 
* Fix eval in trace bugs

* comment nit
This commit is contained in:
Awni Hannun
2024-02-02 09:57:12 -08:00
committed by GitHub
parent 506d43035c
commit cb6156d35d
4 changed files with 33 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
mlx/backend/metal/metal.cpp
View File

@@ -63,7 +63,15 @@ std::function<void()> make_task(
auto s = arr.primitive().stream();
auto command_buffer = increment_command_buffer(s);
auto outputs = arr.outputs();
arr.primitive().eval_gpu(arr.inputs(), outputs);
{
// If the array is a tracer hold a reference
// to its inputs so they don't get donated
std::vector<array> inputs;
if (arr.is_tracer()) {
inputs = arr.inputs();
}
arr.primitive().eval_gpu(arr.inputs(), outputs);
}
std::vector<std::shared_ptr<array::Data>> buffers;
for (auto& in : arr.inputs()) {
buffers.push_back(in.data_shared_ptr());