From 2fde2d2914c5c63abeb4312b8abd328d66a05c65 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 26 Jul 2024 18:41:52 -0700 Subject: [PATCH] hotfix: external_test_speed_theoretical works on 24GB --- test/external/external_test_speed_theoretical.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/external/external_test_speed_theoretical.py b/test/external/external_test_speed_theoretical.py index 8e2efab57d..fd2f825265 100644 --- a/test/external/external_test_speed_theoretical.py +++ b/test/external/external_test_speed_theoretical.py @@ -13,10 +13,10 @@ def _test(tcount, fxn, szmax): ts = [(x+1).realize() for x in ts] Device.default.synchronize() st = time.perf_counter() - out = jfxn(*ts) + out_nbytes = jfxn(*ts).nbytes() Device.default.synchronize() tms.append(time.perf_counter() - st) - gbs = (out.nbytes()+sum(x.nbytes() for x in ts))*1e-9/min(tms) + gbs = (out_nbytes+sum(x.nbytes() for x in ts))*1e-9/min(tms) print(f"{ts[0].nbytes()/(1024*1024):10.0f} MB, {min(tms)*1e3:6.2f} ms GB/s {gbs:<10.2f} {str(ts[0].shape):20s}") allgbs.append(gbs) return max(allgbs)