Update keras_test.py

EnzymeAD · Nov 11, 2024 · 4f3ca1d · 4f3ca1d
1 parent c59309f
commit 4f3ca1d
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/test/keras_test.py b/test/keras_test.py
@@ -87,6 +87,11 @@ def main(argv):
             for name, pipe, dev in pipelines():
                 if pipe is None and AD:
                     continue
+                if bname.endswith("predict") and AD:
+                    continue
+                # 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: