Internal change

PiperOrigin-RevId: 402372196 Change-Id: If7ecd1702e29bf2f0171386581e0ae9f5cd6e4ff
Marbck · Oct 11, 2021 · 5a64f0b · 5a64f0b
1 parent ae3c707
commit 5a64f0b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py
@@ -1175,7 +1175,8 @@ def decorator(cls):
 
     for name, value in cls.__dict__.copy().items():
       if (callable(value) and
-          name.startswith(unittest.TestLoader.testMethodPrefix) and
+          (name.startswith(unittest.TestLoader.testMethodPrefix) or
+           name.startswith("benchmark")) and
           not getattr(value, "_disable_eager_op_as_function", False)):
         setattr(cls, name + "WithEagerOpAsFunctionEnabled",
                 enable_eager_op_as_function(value))