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

Use psutil to get memory info on Windows (#1700)

Browse Source
This commit is contained in:
Cheng 2024-12-14 12:50:13 +09:00 committed by GitHub
parent 635117c5d4
commit dfccd17ab9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 12 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
python/tests/test_array.py
View File

@ -2,14 +2,20 @@
import gc
import operator
import os
import pickle
import resource
import platform
import sys
import unittest
import weakref
from copy import copy, deepcopy
from itertools import permutations
if platform.system() == "Windows":
import psutil
else:
import resource
import mlx.core as mx
import mlx_tests
import numpy as np
@ -1932,7 +1938,11 @@ class TestArray(mlx_tests.MLXTestCase):
def test_siblings_without_eval(self):
def get_mem():
return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
if platform.system() == "Windows":
process = psutil.Process(os.getpid())
return process.memory_info().peak_wset
else:
return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
key = mx.array([1, 2])