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

12
python/tests/test_array.py
View File

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