diff --git a/deepmatch/utils.py b/deepmatch/utils.py
index c52e6d7..81f7664 100644
--- a/deepmatch/utils.py
+++ b/deepmatch/utils.py
@@ -51,7 +51,7 @@ def l2_normalize(x, axis=-1):
 
 
 def inner_product(x, y, temperature=1.0):
-    return Lambda(lambda x: tf.reduce_sum(tf.multiply(x[0], x[1])) / temperature)([x, y])
+    return Lambda(lambda x: tf.reduce_sum(tf.multiply(x[0], x[1]), axis=-1, keepdims=True) / temperature)([x, y])
 
 
 def recall_N(y_true, y_pred, N=50):