A quick fix for a performance regression

PiperOrigin-RevId: 554537362 Change-Id: Id4a8f472f59b52ae86a1587deeec0d577fee3013
google · Aug 7, 2023 · a9b7628 · a9b7628
1 parent 9f48610
commit a9b7628
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/qkeras/utils.py b/qkeras/utils.py
@@ -466,11 +466,8 @@ def get_y_from_TFOpLambda(model_cfg, layer):
 
     # TODO(lishanok): Extend support for other TFOpLambda types when needed
     if op_name == layer.name and  class_name == "TFOpLambda":
-      assert (
-          "tf.__operators__.add" in op_name
-          or "tf.math.add" in op_name
-          or "tf.math.multiply" in op_name
-      ), "TFOpLambda layer {} not supported!".format(op_name)
+      assert ("tf.__operators__.add" in op_name or "tf.math.multiply"
+              in op_name), "TFOpLambda layer {} not supported!".format(op_name)
       return layer_config["inbound_nodes"][-1][-1]["y"]
 
   return None