From 298514d1378f32a69b01967cddb4dd68a5ab898c Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Mon, 28 Jul 2025 12:31:19 -0500 Subject: [PATCH 1/2] Try enabling segment anything --- test/keras_test.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/keras_test.py b/test/keras_test.py index b5d90c0c4..8c3b89865 100644 --- a/test/keras_test.py +++ b/test/keras_test.py @@ -77,8 +77,7 @@ def main(argv): benchfns += [ ("bert_fit", benchmark.bert.bert_fit_run, Both), ] - # also oom - if False: + if True: benchfns += [ ("sam_predict", benchmark.sam.sam_predict_run, Both), ("sam_fit", benchmark.sam.sam_fit_run, Both), From 796579cc93af3dda2cd3bbbb192bb4f83fd307ed Mon Sep 17 00:00:00 2001 From: William Moses Date: Mon, 28 Jul 2025 14:41:38 -0500 Subject: [PATCH 2/2] Update keras_test.py --- test/keras_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/keras_test.py b/test/keras_test.py index 8c3b89865..ef8b9a61e 100644 --- a/test/keras_test.py +++ b/test/keras_test.py @@ -97,6 +97,10 @@ def main(argv): if bname == "bert_predict" and name == "IPartOpt": continue + # Oom's the gpu ci + if bname == "sam_predict" and name == "JaXPipe": + continue + os.environ.pop("ENZYME_JAX", None) os.environ.pop("ENZYME_JAX_PRE", None) if pipe is not None: