Update keras_test.py

EnzymeAD · Nov 11, 2024 · 2154129 · 2154129
1 parent 4f3ca1d
commit 2154129
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/test/keras_test.py b/test/keras_test.py
@@ -70,6 +70,11 @@ def main(argv):
     if True:
         benchfns += [
             ("bert_predict", benchmark.bert.bert_predict_run, Both),
+        ]
+
+    # OOM's the 4070 :'(
+    if False:
+        benchfns += [
             ("bert_fit", benchmark.bert.bert_fit_run, Both),
         ]
 
@@ -92,6 +97,7 @@ def main(argv):
                 # Oom's the gpu ci
                 if bname == "bert_predict" and name == "IPartOpt":
                     continue
+
                 os.environ.pop("ENZYME_JAX", None)
                 os.environ.pop("ENZYME_JAX_PRE", None)
                 if pipe is not None: