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

Make sure that arrays are freed when saving (#247)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2023-12-21 14:08:24 -08:00
committed by GitHub
parent b3916cbf2b
commit 2c7df6795e
3 changed files with 17 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/src/ops.cpp
View File

@@ -2857,18 +2857,20 @@ void init_ops(py::module_& m) {
"file"_a,
"arr"_a,
py::pos_only(),
"retain_graph"_a = true,
"retain_graph"_a = std::nullopt,
py::kw_only(),
R"pbdoc(
save(file: str, arr: array, / , retain_graph: bool = True)
save(file: str, arr: array, / , retain_graph: Optional[bool] = None)
Save the array to a binary file in ``.npy`` format.
Args:
file (str): File to which the array is saved
arr (array): Array to be saved.
retain_graph(bool): Optional argument to retain graph
during array evaluation before saving. Default: True
retain_graph (bool, optional): Optional argument to retain graph
during array evaluation before saving. If not provided the graph
is retained if we are during a function transformation. Default:
None
)pbdoc");
m.def(